[Prompt Optimization Backend PR #1] Wrap prompt optimize in mlflow job by chenmoneygithub · Pull Request #20001 · mlflow/mlflow

chenmoneygithub · 2026-01-15T01:39:02Z

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Wrap the mlflow.genai.optimize_prompts() in an mlflow job, which is a required step for building the mlflow prompt optimization backend.

This is part of #19926, and I am just splitting the PR for easy review.

Note: This PR is safe to go with either 3.9 or 3.10 release, since it's not directly user-facing, but just a backbone for the optimization backend.

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-15T01:39:16Z

@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 wraps the mlflow.genai.optimize_prompts() function in an MLflow job to enable asynchronous prompt optimization execution. This is a foundational component for building the MLflow prompt optimization backend.

Changes:

Added optimize_prompts_job function as a background job wrapper
Implemented helper functions for optimizer creation, scorer loading, and predict function building
Added telemetry event tracking for job execution
Registered the new job in the MLflow server's job registry

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/job.py	New module implementing the prompt optimization job wrapper with helper functions for optimizer creation, scorer loading, and predict function building
mlflow/telemetry/events.py	Added OptimizePromptsJobEvent for tracking job execution telemetry
mlflow/server/jobs/init.py	Registered the optimize_prompts_job in the supported job function list and allowed job name list
tests/genai/optimize/test_job.py	Comprehensive unit tests for the job wrapper and helper functions
tests/telemetry/test_events.py	Unit tests for the OptimizePromptsJobEvent telemetry event

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

mlflow/genai/optimize/job.py

tests/genai/optimize/test_job.py

mlflow/genai/optimize/job.py

github-actions · 2026-01-15T01:46:45Z

Documentation preview for 13ca846 is available at:

https://pr-20001--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-15T07:24:14Z

mlflow/genai/optimize/job.py

+        Dict containing optimization results and metadata.
+    """
+    # Record telemetry event for job execution
+    _record_event(


can we use the decorator style for adding telemetry for this api?

TomeHirata · 2026-01-15T07:25:24Z

mlflow/genai/optimize/job.py

+    prompt_uri: str,
+    dataset_id: str,
+    optimizer_type: str,
+    optimizer_config_json: str | None,


Does it cause any issue if we accept optimizer_config as dict[str, Any] instead of str in this method?

oops, this is supposed to be a dict, and handler code should deserialize the string to dict. updated

TomeHirata · 2026-01-15T07:27:14Z

mlflow/genai/optimize/job.py

+    dataset_id: str,
+    optimizer_type: str,
+    optimizer_config_json: str | None,
+    scorers: list[str],


nit: scorer_names?

Q: What's the format for custom scorers ?

TomeHirata · 2026-01-15T07:28:53Z

mlflow/genai/optimize/job.py

+        dataset_id: The ID of the EvaluationDataset containing training data.
+        optimizer_type: The optimizer type string (e.g., "gepa", "metaprompt").
+        optimizer_config_json: JSON string of optimizer-specific configuration.
+        scorers: List of scorer names. Can be built-in scorer class names


What would be the workflow if users want to use a custom llm judge? Is it like create a judge using make_judge -> register it using .register() -> call this api with its judge name?

yes exactly

TomeHirata · 2026-01-15T07:35:42Z

mlflow/genai/optimize/job.py

+    source_prompt = load_prompt(prompt_uri)
+
+    # Resume the given run ID. Params have already been logged by the handler
+    with start_run(run_id=run_id):


Does it mean a MLflow run is created and ended before this method is called?

yes, it's written this way because we need to associate the optimization job with an mlflow run when the job starts, if we don't create the run before kicking off the job, we lose this lineage until finishing the job.

Code in handler should be more straightforward

TomeHirata · 2026-01-15T07:36:24Z

mlflow/genai/optimize/job.py

+    config = json.loads(optimizer_config_json) if optimizer_config_json else {}
+    optimizer_type = optimizer_type.lower() if optimizer_type else ""
+
+    if optimizer_type == "gepa":


Can we define an enum for optimizer type?

TomeHirata · 2026-01-15T07:36:51Z

mlflow/genai/optimize/job.py

+    Returns:
+        A callable that takes inputs dict and returns the LLM response.
+    """
+    import litellm


can we display a kind message if litellm is not installed?

TomeHirata · 2026-01-15T07:39:42Z

mlflow/genai/optimize/job.py

+    set_experiment(experiment_id=experiment_id)
+
+    dataset = get_dataset(dataset_id=dataset_id)
+    train_data = dataset.to_df()


nit: I guess optimize_prompts accepts EvaluationDataset

WeichenXu123 · 2026-01-15T08:42:20Z

mlflow/genai/optimize/job.py

+            the experiment's scorer registry.
+
+    Returns:
+        Dict containing optimization results and metadata.


Reminder: need to guarantee that the dict is JSON serializable.

chenmoneygithub · 2026-01-16T04:05:01Z

mlflow/genai/optimize/job.py

+    # Resume the given run ID. Params have already been logged by the handler
+    with start_run(run_id=run_id):
+        # Link source prompt to run for lineage
+        client = MlflowClient()


discussed offline, this should be fine since we expect the job to be kicked off by the server handler

TomeHirata · 2026-01-16T10:14:46Z

mlflow/genai/optimize/job.py

+            enable_tracking=True,
+        )
+
+    return {


nit: shall we deflne a class for the response?

good call, done!

TomeHirata

LGTM

wrap prompt optimize in mlflow job

07f471a

Copilot AI review requested due to automatic review settings January 15, 2026 01:39

github-actions bot added the area/prompts MLflow Prompt Registry and Optimization label Jan 15, 2026

Copilot started reviewing on behalf of chenmoneygithub January 15, 2026 01:39 View session

github-actions bot added the rn/feature Mention under Features in Changelogs. label Jan 15, 2026

Copilot AI reviewed Jan 15, 2026

View reviewed changes

mlflow/genai/optimize/job.py Outdated Show resolved Hide resolved

mlflow/genai/optimize/job.py Outdated Show resolved Hide resolved

tests/genai/optimize/test_job.py Show resolved Hide resolved

mlflow/genai/optimize/job.py Outdated Show resolved Hide resolved

chenmoneygithub added 2 commits January 14, 2026 17:56

clean up

66b4a42

fix copilot comments

4a8d476

chenmoneygithub requested review from TomeHirata and WeichenXu123 January 15, 2026 06:32

TomeHirata reviewed Jan 15, 2026

View reviewed changes

WeichenXu123 reviewed Jan 15, 2026

View reviewed changes

github-actions bot assigned TomeHirata and WeichenXu123 Jan 15, 2026

WeichenXu123 reviewed Jan 15, 2026

View reviewed changes

chenmoneygithub added 3 commits January 15, 2026 20:00

address comments

3208a36

clean up

a78379f

Merge branch 'master' into mlflow-po-backend-pr-1

1379a31

TomeHirata reviewed Jan 16, 2026

View reviewed changes

TomeHirata approved these changes Jan 16, 2026

View reviewed changes

use a dataclass to hold the result

13ca846

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

Merged via the queue into mlflow:master with commit 6b2e131 Jan 16, 2026
47 of 48 checks passed

chenmoneygithub deleted the mlflow-po-backend-pr-1 branch January 16, 2026 20:48

Conversation

chenmoneygithub commented Jan 15, 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 15, 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 15, 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!

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!

This comment was marked as outdated.

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!

TomeHirata left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

chenmoneygithub commented Jan 15, 2026 •

edited

Loading

github-actions bot commented Jan 15, 2026 •

edited

Loading