Fix code example in MLFlow section of deployment docs (#8229)

* fix(documentation): mlflow deployment code example

* add caveat callout

---------

Co-authored-by: chenmoneygithub <chen.qian@databricks.com>

Author: jmho

docs/docs/tutorials/deployment/index.md

@@ -120,7 +120,13 @@ print(response.json())
 You should see the response like below:
 
 ```json
-{'status': 'success', 'data': {'reasoning': 'The capital of France is a well-known fact, commonly taught in geography classes and referenced in various contexts. Paris is recognized globally as the capital city, serving as the political, cultural, and economic center of the country.', 'answer': 'The capital of France is Paris.'}}
+{
+  "status": "success",
+  "data": {
+    "reasoning": "The capital of France is a well-known fact, commonly taught in geography classes and referenced in various contexts. Paris is recognized globally as the capital city, serving as the political, cultural, and economic center of the country.",
+    "answer": "The capital of France is Paris."
+  }
+}
 ```
 
 ## Deploying with MLflow
@@ -140,7 +146,14 @@ Let's spin up the MLflow tracking server, where we will store our DSPy program.
 ```
 
 Then we can define the DSPy program and log it to the MLflow server. "log" is an overloaded term in MLflow, basically it means
-we store the program information along with environment requirements in the MLflow server. See the code below:
+we store the program information along with environment requirements in the MLflow server. This is done via the `mlflow.dspy.log_model()`
+function, please see the code below:
+
+> [!NOTE]
+> As of MLflow 2.22.0, there is a caveat that you must wrap your DSPy program in a custom DSPy Module class when deploying with MLflow.
+> This is because MLflow requires positional arguments while DSPy pre-built modules disallow positional arguments, e.g., `dspy.Predict`
+> or `dspy.ChainOfThought`. To work around this, create a wrapper class that inherits from `dspy.Module` and implement your program's
+> logic in the `forward()` method, as shown in the example below.
 
 ```python
 import dspy
@@ -151,7 +164,16 @@ mlflow.set_experiment("deploy_dspy_program")
 
 lm = dspy.LM("openai/gpt-4o-mini")
 dspy.settings.configure(lm=lm)
-dspy_program = dspy.ChainOfThought("question -> answer")
+
+class MyProgram(dspy.Module):
+    def __init__(self):
+        super().__init__()
+        self.cot = dspy.ChainOfThought("question -> answer")
+
+    def forward(self, question):
+        return self.cot(question=question)
+
+dspy_program = MyProgram()
 
 with mlflow.start_run():
     mlflow.dspy.log_model(
@@ -188,7 +210,18 @@ After the program is deployed, you can test it with the following command:
 You should see the response like below:
 
 ```json
-{"choices": [{"index": 0, "message": {"role": "assistant", "content": "{\"reasoning\": \"The question asks for the sum of 2 and 2. To find the answer, we simply add the two numbers together: 2 + 2 = 4.\", \"answer\": \"4\"}"}, "finish_reason": "stop"}]}
+{
+  "choices": [
+    {
+      "index": 0,
+      "message": {
+        "role": "assistant",
+        "content": "{\"reasoning\": \"The question asks for the sum of 2 and 2. To find the answer, we simply add the two numbers together: 2 + 2 = 4.\", \"answer\": \"4\"}"
+      },
+      "finish_reason": "stop"
+    }
+  ]
+}
 ```
 
 For complete guide on how to deploy a DSPy program with MLflow, and how to customize the deployment, please refer to the