|
| 1 | +# Tracking DSPy Optimizers with MLflow |
| 2 | + |
| 3 | +This tutorial demonstrates how to use MLflow to track and analyze your DSPy optimization process. MLflow's built-in integration for DSPy provides traceability and debuggability for your DSPy optimization experience. It allows you to understand the intermediate trials during the optimization, store the optimized program and its results, and provides observability into your program execution. |
| 4 | + |
| 5 | +Through the autologging capability, MLflow tracks the following information: |
| 6 | + |
| 7 | +* **Optimizer Parameters** |
| 8 | + * Number of few-shot examples |
| 9 | + * Number of candidates |
| 10 | + * Other configuration settings |
| 11 | + |
| 12 | +* **Program States** |
| 13 | + * Initial instuctions and few-shot examples |
| 14 | + * Optimized instuctions and few-shot examples |
| 15 | + * Intermediate instuctions and few-shot examples during optimization |
| 16 | + |
| 17 | +* **Datasets** |
| 18 | + * Training data used |
| 19 | + * Evaluation data used |
| 20 | + |
| 21 | +* **Performance Progression** |
| 22 | + * Overall metric progression |
| 23 | + * Performance at each evaluation step |
| 24 | + |
| 25 | +* **Traces** |
| 26 | + * Program execution traces |
| 27 | + * Model responses |
| 28 | + * Intermediate prompts |
| 29 | + |
| 30 | +## Getting Started |
| 31 | + |
| 32 | +### 1. Install MLflow |
| 33 | +First, install MLflow (version 2.21.1 or later): |
| 34 | + |
| 35 | +```bash |
| 36 | +pip install mlflow>=2.21.1 |
| 37 | +``` |
| 38 | + |
| 39 | +### 2. Start MLflow Tracking Server |
| 40 | + |
| 41 | +Let's spin up the MLflow tracking server with the following command. This will start a local server at `http://127.0.0.1:5000/`: |
| 42 | + |
| 43 | +```bash |
| 44 | +# It is highly recommended to use SQL store when using MLflow tracing |
| 45 | +mlflow server --backend-store-uri sqlite:///mydb.sqlite |
| 46 | +``` |
| 47 | + |
| 48 | +### 3. Enable Autologging |
| 49 | + |
| 50 | +Configure MLflow to track your DSPy optimization: |
| 51 | + |
| 52 | +```python |
| 53 | +import mlflow |
| 54 | +import dspy |
| 55 | + |
| 56 | +# Enable autologging with all features |
| 57 | +mlflow.dspy.autolog( |
| 58 | + log_compiles=True, # Track optimization process |
| 59 | + log_evals=True, # Track evaluation results |
| 60 | + log_traces_from_compile=True # Track program traces during optimization |
| 61 | +) |
| 62 | + |
| 63 | +# Configure MLflow tracking |
| 64 | +mlflow.set_tracking_uri("http://localhost:5000") # Use local MLflow server |
| 65 | +mlflow.set_experiment("DSPy-Optimization") |
| 66 | +``` |
| 67 | + |
| 68 | +### 4. Optimizing Your Program |
| 69 | + |
| 70 | +Here's a complete example showing how to track the optimization of a math problem solver: |
| 71 | + |
| 72 | +```python |
| 73 | +import dspy |
| 74 | +from dspy.datasets.gsm8k import GSM8K, gsm8k_metric |
| 75 | + |
| 76 | +# Configure your language model |
| 77 | +lm = dspy.LM(model="openai/gpt-4o") |
| 78 | +dspy.configure(lm=lm) |
| 79 | + |
| 80 | +# Load dataset |
| 81 | +gsm8k = GSM8K() |
| 82 | +trainset, devset = gsm8k.train, gsm8k.dev |
| 83 | + |
| 84 | +# Define your program |
| 85 | +program = dspy.ChainOfThought("question -> answer") |
| 86 | + |
| 87 | +# Create and run optimizer with tracking |
| 88 | +teleprompter = dspy.teleprompt.MIPROv2( |
| 89 | + metric=gsm8k_metric, |
| 90 | + auto="light", |
| 91 | +) |
| 92 | + |
| 93 | +# The optimization process will be automatically tracked |
| 94 | +optimized_program = teleprompter.compile( |
| 95 | + program, |
| 96 | + trainset=trainset, |
| 97 | +) |
| 98 | +``` |
| 99 | + |
| 100 | +### 5. Viewing Results |
| 101 | + |
| 102 | +Once your optimization is complete, you can analyze the results through MLflow's UI. Let's walk through how to explore your optimization runs. |
| 103 | + |
| 104 | +#### Step 1: Access the MLflow UI |
| 105 | +Navigate to `http://localhost:5000` in your web browser to access the MLflow tracking server UI. |
| 106 | + |
| 107 | +#### Step 2: Understanding the Experiment Structure |
| 108 | +When you open the experiment page, you'll see a hierarchical view of your optimization process. The parent run represents your overall optimization process, while the child runs show each intermediate version of your program that was created during optimization. |
| 109 | + |
| 110 | + |
| 111 | + |
| 112 | +#### Step 3: Analyzing the Parent Run |
| 113 | +Clicking on the parent run reveals the big picture of your optimization process. You'll find detailed information about your optimizer's configuration parameters and how your evaluation metrics progressed over time. The parent run also stores your final optimized program, including the instructions, signature definitions, and few-shot examples that were used. Additionally, you can review the training data that was used during the optimization process. |
| 114 | + |
| 115 | + |
| 116 | + |
| 117 | +#### Step 4: Examining Child Runs |
| 118 | +Each child run provides a detailed snapshot of a specific optimization attempt. When you select a child run from the experiment page, you can explore several aspects of that particular intermediate program. |
| 119 | +On the run parameter tab or artifact tab, you can review the instructions and few-shot examples used for the intermediate program. |
| 120 | +One of the most powerful features is the Traces tab, which provides a step-by-step view of your program's execution. Here you can understand exactly how your DSPy program processes inputs and generates outputs. |
| 121 | + |
| 122 | + |
| 123 | + |
| 124 | +### 6. Loading Models for Inference |
| 125 | +You can load the optimized program directly from the MLflow tracking server for inference: |
| 126 | + |
| 127 | +```python |
| 128 | +model_path = mlflow.artifacts.download_artifacts("mlflow-artifacts:/path/to/best_model.json") |
| 129 | +program.load(model_path) |
| 130 | +``` |
| 131 | + |
| 132 | +## Troubleshooting |
| 133 | + |
| 134 | +- If traces aren't appearing, ensure `log_traces_from_compile=True` |
| 135 | +- For large datasets, consider setting `log_traces_from_compile=False` to avoid memory issues |
| 136 | +- Use `mlflow.get_run(run_id)` to programmatically access MLflow run data |
| 137 | + |
| 138 | +For more features, explore the [MLflow Documentation](https://mlflow.org/docs/latest/llms/dspy). |
0 commit comments