Track intermediate candidates and evaluation scores in gepa optimizer by chenmoneygithub · Pull Request #20198 · mlflow/mlflow

chenmoneygithub · 2026-01-21T21:59:51Z

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Track intermediate candidates and evaluation scores in gepa optimizer, which is useful to track the GEPA optimization process. See the screenshot below for a sample:

We are tracking 3 things for each candidate that hits the full validation:

The overall score (aggregated across each data)
The eval table for each data record
The candidate prompt. In the case where multiple prompts are involved, each prompt will be a separate file for better readability.

Please note that we don't track the candidate that doesn't hit full validation, which is not addd to the pareto frontier of GEPA optimizer.

How is this PR tested?

Existing unit/integration tests
New unit/integration tests
Manual tests

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

No. You can skip the rest of this section.
Yes. Give a description of this change to be included in the release notes for MLflow users.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

How should the PR be classified in the release notes? Choose one:

rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
rn/feature - A new user-facing feature worth mentioning in the release notes
rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
rn/documentation - A user-facing documentation change worth mentioning in the release notes

Should this PR be included in the next patch release?

Yes should be selected for bug fixes, documentation updates, and other small changes. No should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

What is a minor/patch release?

Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
Bug fixes, doc updates and new features usually go into minor releases.
Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
Bug fixes and doc updates usually go into patch releases.

Yes (this PR will be cherry-picked and included in the next patch release)
No (this PR will be included in the next minor release)

github-actions · 2026-01-21T22:00:11Z

🛠 DevTools 🛠

Install mlflow from this PR

# mlflow
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/20198/merge
# mlflow-skinny
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/20198/merge#subdirectory=libs/skinny

For Databricks, use the following command:

%sh curl -LsSf https://raw.githubusercontent.com/mlflow/mlflow/HEAD/dev/install-skinny.sh | sh -s pull/20198/merge

github-actions · 2026-01-21T22:00:12Z

@chenmoneygithub Thank you for the contribution! Could you fix the following issue(s)?

⚠ DCO check

The DCO check failed. Please sign off your commit(s) by following the instructions here. See https://github.com/mlflow/mlflow/blob/master/CONTRIBUTING.md#sign-your-work for more details.

Copilot

Pull request overview

This PR adds tracking capabilities to the GEPA optimizer to monitor intermediate candidates and their evaluation scores during the optimization process. This enhancement provides visibility into the optimization workflow by logging validation candidates, aggregate scores, per-record scores, and individual scorer results.

Changes:

Extended EvaluationResultRecord to include individual scorer results
Modified evaluation metric functions to return individual scores alongside aggregate scores
Implemented artifact logging for validation candidates in the GEPA optimizer
Added comprehensive test coverage for the new tracking functionality

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
mlflow/genai/optimize/types.py	Added `individual_scores` field to `EvaluationResultRecord` with default factory
mlflow/genai/optimize/util.py	Updated `create_metric_from_scorers` to return tuple with individual scores
mlflow/genai/optimize/optimize.py	Modified `_build_eval_fn` to unpack and pass individual scores from metric function
mlflow/genai/optimize/optimizers/gepa_optimizer.py	Added tracking logic in `MlflowGEPAAdapter` with `_log_validation_candidate` method for artifact logging
tests/genai/optimize/optimizers/test_gepa_optimizer.py	Added comprehensive test for prompt candidate logging functionality

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

mlflow/genai/optimize/optimizers/gepa_optimizer.py

github-actions · 2026-01-21T22:10:52Z

Documentation preview for 9f524f0 is available at:

https://pr-20198--mlflow-docs-preview.netlify.app/docs/latest/

More info

Ignore this comment if this PR does not change the documentation.
The preview is updated when a new commit is pushed to this PR.
This comment was created by this workflow run.
The documentation was built by this workflow run.

TomeHirata · 2026-01-26T04:41:56Z

mlflow/genai/optimize/util.py

+            # Log per-scorer metrics (for API response parsing)
+            if output.initial_eval_score_per_scorer:
+                for scorer_name, score in output.initial_eval_score_per_scorer.items():
+                    mlflow.log_metric(f"initial_eval_score.{scorer_name}", score)


Let's use log_metrics for efficient logging

TomeHirata · 2026-01-26T04:45:38Z

mlflow/genai/optimize/optimizers/gepa_optimizer.py

                    batch: List of data instances to evaluate
                    candidate: Proposed text components (prompts)
-                    capture_traces: Whether to capture execution traces
+                    capture_traces: Whether to capture execution traces.


Can we not add this period or add periods to all args for consistency?

TomeHirata · 2026-01-26T04:48:35Z

mlflow/genai/optimize/optimizers/gepa_optimizer.py

                outputs = [result.outputs for result in eval_results]
                scores = [result.score for result in eval_results]
                trajectories = eval_results if capture_traces else None
+                objective_scores = [result.individual_scores for result in eval_results]


q: is this parameter present for all supported versions?

I will bump the GEPA requirement to be >=0.0.26, there was a critical bug fixed in this PR gepa-ai/gepa#171, which has been released with 0.0.26. So yes, the parameter is available for valid gepa versions.

TomeHirata · 2026-01-26T04:51:26Z

mlflow/genai/optimize/optimizers/gepa_optimizer.py

+                    candidate: The candidate prompts being validated
+                    eval_results: Evaluation results containing scores
+                """
+                import mlflow


nit: Can we move this to the module level?

TomeHirata · 2026-01-26T04:51:57Z

mlflow/genai/optimize/optimizers/gepa_optimizer.py

+                """
+                import mlflow
+
+                active_run = mlflow.active_run()


any reason not to use self.tracking_enabled?

good call, changed

TomeHirata · 2026-01-26T04:54:02Z

mlflow/genai/optimize/optimizers/gepa_optimizer.py

+                # Compute per-scorer average scores
+                scorer_names = set()
+                for result in eval_results:
+                    scorer_names.update(result.individual_scores.keys())


nit: we can use scorer_names |= result.individual_scores.keys()

gotcha, changed

TomeHirata · 2026-01-26T04:58:19Z

mlflow/genai/optimize/optimizers/gepa_optimizer.py

+
+                # Build the evaluation results table and log to MLflow as a table artifact
+                eval_results_table = {
+                    "inputs": [json.dumps(r.inputs) for r in eval_results],


Do we need json.dumps here? isn't this handled by mlflow.log_table?

good call, I wasn't aware of that part, removed!

TomeHirata · 2026-01-26T04:59:19Z

mlflow/genai/optimize/optimizers/gepa_optimizer.py

+                for result in eval_results:
+                    scorer_names.update(result.individual_scores.keys())
+
+                per_scorer_scores = {}


Can we move this logic below so that the variable generation logic is closer to where is't used?

TomeHirata · 2026-01-26T05:00:34Z

mlflow/genai/optimize/optimizers/gepa_optimizer.py

+                        prompt_path = tmp_path / f"{prompt_name}.txt"
+                        with open(prompt_path, "w") as f:
+                            f.write(prompt_text)
+                        mlflow.log_artifact(str(prompt_path), artifact_path=iteration_dir)


Can't we pass Path to mlflow.log_artifact?

good call, AI-redundant code, changed!

TomeHirata · 2026-01-26T05:01:51Z

mlflow/genai/optimize/optimizers/gepa_optimizer.py

+                        r.individual_scores.get(scorer_name) for r in eval_results
+                    ]
+
+                iteration_dir = f"prompt_candidates/iteration_{iteration}"


I guess these path names will be used for returning the optimization result from the rest API, so shall we extract these path/file names as consts?

good call, the new commit defines a few constants for reusing.

mlflow/genai/optimize/optimizers/gepa_optimizer.py

TomeHirata · 2026-01-26T05:09:41Z

mlflow/genai/optimize/types.py

    score: float | None
    trace: Trace
    rationales: dict[str, str]
+    individual_scores: dict[str, float] = field(default_factory=dict)


Per scorer score is not necessarily numeric. Please refer to Scorer.__call__

Yes that's a good callout, I am now planning to only allow numeric scorers: #20197 (comment)

let's discuss!

TomeHirata · 2026-01-26T05:10:45Z

mlflow/genai/optimize/optimizers/gepa_optimizer.py

+                        if scorer_name in r.individual_scores
+                    ]
+                    if scores:
+                        per_scorer_scores[scorer_name] = sum(scores) / len(scores)


I think individual scorers cannot always be aggregated using sum. Please see my comment below.

similar to #20198 (comment).

mlflow/genai/optimize/optimizers/gepa_optimizer.py

TomeHirata · 2026-01-26T05:13:34Z

mlflow/genai/optimize/optimizers/gepa_optimizer.py

+                # Log scores summary as JSON artifact
+                scores_data = {
+                    "aggregate": aggregate_score,
+                    "per_scorer": per_scorer_scores,


I wonder if we should log metrics for each scorer name where possible. This enables users to see the time progression of each scorer result.

yes good idea, done!

TomeHirata

LGTM if we support string output in the follow up PR

mlflow/genai/optimize/optimizers/gepa_optimizer.py

track intermediate candidates in gepa optimizer

f3e3810

Copilot AI review requested due to automatic review settings January 21, 2026 21:59

Copilot started reviewing on behalf of chenmoneygithub January 21, 2026 22:00 View session

github-actions bot added area/prompts MLflow Prompt Registry and Optimization rn/feature Mention under Features in Changelogs. and removed rn/feature Mention under Features in Changelogs. area/prompts MLflow Prompt Registry and Optimization labels Jan 21, 2026

formatting

358640b

Copilot AI reviewed Jan 21, 2026

View reviewed changes

chenmoneygithub added 2 commits January 21, 2026 15:22

fix copilot comments

282fa44

merge master

150c4de

github-actions bot added area/prompts MLflow Prompt Registry and Optimization rn/feature Mention under Features in Changelogs. labels Jan 23, 2026

chenmoneygithub requested a review from TomeHirata January 23, 2026 02:27

chenmoneygithub added 2 commits January 23, 2026 13:50

properly handle individual scorer scores

5d89369

refine code

ae3fe30

TomeHirata reviewed Jan 26, 2026

View reviewed changes

mlflow/genai/optimize/optimizers/gepa_optimizer.py Outdated Show resolved Hide resolved

TomeHirata reviewed Jan 26, 2026

View reviewed changes

mlflow/genai/optimize/optimizers/gepa_optimizer.py Show resolved Hide resolved

TomeHirata reviewed Jan 26, 2026

View reviewed changes

github-actions bot assigned TomeHirata Jan 26, 2026

fix comments

c65d4fb

chenmoneygithub requested a review from TomeHirata January 26, 2026 20:55

Merge branch 'master' into track-gepa-intermediate-prompts

9f524f0

TomeHirata approved these changes Jan 27, 2026

View reviewed changes

chenmoneygithub added this pull request to the merge queue Jan 27, 2026

Merged via the queue into mlflow:master with commit dfdc877 Jan 27, 2026
47 of 50 checks passed

chenmoneygithub deleted the track-gepa-intermediate-prompts branch January 27, 2026 06:13

TomeHirata reviewed Jan 28, 2026

View reviewed changes

mlflow/genai/optimize/optimizers/gepa_optimizer.py Show resolved Hide resolved

Conversation

chenmoneygithub commented Jan 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Related Issues/PRs

What changes are proposed in this pull request?

How is this PR tested?

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

What component(s), interfaces, languages, and integrations does this PR affect?

How should the PR be classified in the release notes? Choose one:

Should this PR be included in the next patch release?

Uh oh!

github-actions bot commented Jan 21, 2026

Install mlflow from this PR

Uh oh!

github-actions bot commented Jan 21, 2026

⚠ DCO check

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Jan 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

TomeHirata Jan 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chenmoneygithub commented Jan 21, 2026 •

edited

Loading

github-actions bot commented Jan 21, 2026 •

edited

Loading

TomeHirata Jan 26, 2026 •

edited

Loading