[Bug] Issue with dspy and mlflow Deployment - TypeError: ChainOfThought.forward() takes 1 positional argument but 2 were given

[Jammillk]

What happened?

Title: Issue with dspy and mlflow Deployment - TypeError: ChainOfThought.forward() takes 1 positional argument but 2 were given

Description:

Hi, I am following the official DSPy documentation to deploy a DSPy application with MLflow (as outlined in the tutorial). However, I am encountering the following issue despite following the example code almost exactly.

PS：this issue also aks mlflow：mlflow/mlflow#15546

Steps to reproduce

Steps Taken:

I used the example code directly from the DSPy documentation with minor adjustments for my setup (e.g., model name and experiment name), but I received the error below.

Code (Problematic):

import dspy
import mlflow

mlflow.dspy.autolog()

# This is optional. Create an MLflow Experiment to store and organize your traces.
mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("DSPy")

lm = dspy.LM(
    model="openai/deepseek-v3"
    # ...
)
dspy.configure(lm=lm)

dspy_program = dspy.ChainOfThought("question -> answer", n=3)

input_example = {"messages": [{"role": "user", "content": "What is LLM agent?"}]}

with mlflow.start_run():
    mlflow.dspy.log_model(
        dspy_program,
        "dspy_program",
        input_example=input_example,  # Ensure input format is correct
        task="llm/v1/chat"
    )

Error (Exception Stack):

2025/04/29 10:10:16 WARNING mlflow.utils.requirements_utils: Failed to run predict on input_example, dependencies introduced in predict are not captured.
TypeError('ChainOfThought.forward() takes 1 positional argument but 2 were given')Traceback (most recent call last):
  File "D:\devtool\py\lib\site-packages\mlflow\utils\_capture_modules.py", line 166, in load_model_and_predict
    model.predict(input_example, params=params)
  File "D:\devtool\py\lib\site-packages\mlflow\dspy\wrapper.py", line 89, in predict
    outputs = self.model(converted_inputs)
  File "D:\devtool\py\lib\site-packages\dspy\utils\callback.py", line 326, in sync_wrapper
    return fn(instance, *args, **kwargs)
  File "D:\devtool\py\lib\site-packages\dspy\primitives\program.py", line 32, in __call__
    return self.forward(*args, **kwargs)
TypeError: ChainOfThought.forward() takes 1 positional argument but 2 were given
...

Other Example (Working):

I found that a different example, where a forward() method is explicitly defined in a custom class, works correctly. Here is the working example:

import dspy
import mlflow

mlflow.dspy.autolog()

# This is optional. Create an MLflow Experiment to store and organize your traces.
mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("DSPy")

lm = dspy.LM(
    model="openai/deepseek-v3"
    # ...
)
dspy.configure(lm=lm)

# Define a Chain of Thought module
class CoT(dspy.Module):
    def __init__(self):
        super().__init__()
        self.prog = dspy.ChainOfThought("question -> answer")

    def forward(self, question):
        return self.prog(question=question)


dspy_model = CoT()

# Start an MLflow run
with mlflow.start_run():
    # Log the model
    model_info = mlflow.dspy.log_model(
        dspy_model,
        artifact_path="model",
        input_example="what is 2 + 2?",
    )

Result (Working Example):

Downloading artifacts: 100%|██████████| 7/7 [00:00<00:00, 6983.86it/s]
🏃 View run delightful-grub-263 at: http://localhost:5000/#/experiments/381354332525629159/runs/3a29e716bb3e427792e692bbf882aa58
🧪 View experiment at: http://localhost:5000/#/experiments/381354332525629159

Question:

Is this a bug in the DSPy library? I followed the official tutorial exactly but am getting the above error. I noticed that explicitly defining a custom forward() method (as in the second example) works, but the default ChainOfThought seems to break with mlflow.dspy.log_model.

DSPy version

dspy version 2.6.21 mlflow, version 2.22.0

[okhat]

Usually this error is about passing explicit named kwargs

[Jammillk]

Usually this error is about passing explicit named kwargs

Hi,

Thank you for the quick response.

I understand that this error is usually related to passing explicit named keyword arguments. However, I would like to clarify that in my case, I haven't manually passed any named keyword arguments. The code I am running is directly copied from the official DSPy documentation tutorial here, with only minor adjustments for my specific setup (e.g., the model name and experiment name).

Since I followed the official example exactly and encountered this issue, I believe this might be an issue with how the ChainOfThought module is being used with MLflow, as shown in the tutorial.

Could you please help me understand if there is something specific in the code that I might have missed, or if there is any update or fix needed in the DSPy/MLflow integration?

Thanks again!

[Jammillk]

Usually this error is about passing explicit named kwargs

I wanted to bring your attention to the response I received from the MLflow team regarding the issue I encountered. They clarified that the error occurred because I was directly using the raw DSPy ChainOfThought without wrapping it in a custom module. According to them, MLflow expects models to handle inputs like DSPy Modules, where inputs are passed as arguments to the forward() method, not as a dictionary. You can find the details of their response here: MLflow Issue #15546.

Given this, I wanted to check with you whether this approach aligns with DSPy’s intended usage with MLflow. Would it make sense to update the official DSPy tutorial on deployment (https://dspy.ai/tutorials/deployment/#deploying-with-mlflow) to reflect this? The current example might be misleading to users, as it directly uses ChainOfThought without a custom module wrapper, which could lead to confusion.

Looking forward to your thoughts on this.

Thanks!

[okhat]

Ah, thanks @Jammillk ! We should check this. Will discuss with @chenmoneygithub

[chenmoneygithub]

@Jammillk What's the Mlflow version you are using?