Link Prompts and Experiments when prompts are loaded/registered by TomeHirata · Pull Request #18883 · mlflow/mlflow

TomeHirata · 2025-11-18T08:04:46Z

🛠 DevTools 🛠

Install mlflow from this PR

# mlflow
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/18883/merge
# mlflow-skinny
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/18883/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/18883/merge

Related Issues/PRs

n/a

What changes are proposed in this pull request?

Currently linkage to prompt version is tracked for runs, spans, and logged models. In order to allow users to group prompts by experiments, we want to track the linkage between experiments and prompts.

Changes:

Create the linkage between experiments and prompts when load_prompt or register_prompt is called using prompt tags

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)

- Introduced functionality to link prompt versions to experiments, allowing for better organization and tracking of prompts within specific experiments. - Updated the `AbstractStore` and its SQLAlchemy implementation to include methods for adding and retrieving prompt-experiment associations. - Enhanced the `MlflowClient` to facilitate linking prompts during registration and loading processes. - Added tests to ensure correct linking behavior and to verify that prompts can be associated with multiple experiments without duplication. Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

- Updated the prompt linking mechanism to tag prompts with associated experiment IDs instead of creating entity associations. - Modified the `AbstractStore` and its SQLAlchemy implementation to remove methods related to prompt-experiment associations. - Enhanced the `MlflowClient` to manage experiment ID tags during prompt registration and loading. - Updated tests to verify the new tagging behavior and ensure idempotency when tagging prompts with the same experiment ID. Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

…tore class Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

…es and update search_prompts function to return a list instead of PagedList. Enhance docstring for clarity on prompt search functionality. Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

Copilot

Pull Request Overview

This PR adds the ability to link prompts to experiments when prompts are loaded or registered within an active MLflow run. Experiment IDs are stored as comma-separated values in a special prompt tag, enabling users to search for prompts by experiment.

Tracks experiment linkage via PROMPT_EXPERIMENT_IDS_TAG_KEY tag when load_prompt or register_prompt is called within an active run
Implements experiment_id filter syntax for search_prompts() API that translates to tag-based filtering
Adds integration tests covering prompt-experiment linking and search functionality

Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
tests/genai/prompts/test_prompts.py	Adds comprehensive integration tests for prompt-experiment linking and searching by experiment ID
mlflow/tracking/client.py	Updates `register_prompt` to link prompts to experiments and enhances `search_prompts` documentation with experiment_id filter examples
mlflow/tracking/_model_registry/fluent.py	Updates `load_prompt` to link prompts to experiments, improves `search_prompts` documentation, and fixes deprecated `_logger.warn` usage
mlflow/store/model_registry/abstract_store.py	Implements `_parse_experiment_id_filter` to transform experiment_id filter syntax to tag-based LIKE queries
mlflow/store/_unity_catalog/registry/rest_store.py	Integrates `_parse_experiment_id_filter` into Unity Catalog search_prompts implementation
mlflow/prompt/constants.py	Adds `PROMPT_EXPERIMENT_IDS_TAG_KEY` constant for storing experiment IDs in prompt tags
mlflow/entities/entity_type.py	Adds `PROMPT` entity association type for consistency

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

mlflow/store/model_registry/abstract_store.py

mlflow/tracking/client.py

mlflow/tracking/_model_registry/fluent.py

mlflow/store/model_registry/abstract_store.py

mlflow/tracking/client.py

mlflow/tracking/_model_registry/fluent.py

- Updated the prompt linking mechanism to ensure prompts are tagged with associated experiment IDs during registration and loading. - Refactored the `MlflowClient` to manage experiment ID tags more effectively, including asynchronous handling for linking prompts to experiments. - Improved SQL injection protection in the `AbstractStore` by escaping backslashes in experiment IDs. - Updated tests to verify the new tagging behavior and ensure correct linking of prompts to experiments without duplication. Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

github-actions · 2025-11-18T10:05:04Z

Documentation preview for c1d4e23 is available at:

https://pr-18883--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 · 2025-11-18T10:09:10Z

mlflow/tracking/client.py

+                    # Use LIKE to match the experiment ID and experiment ID is auto-incremented
+                    # integer. So add comma before and after the list of experiment IDs to
+                    # avoid false matches (e.g., "1" matches "10").
+                    exp_ids = f",{exp_ids},"


This is not very beautiful, but I couldn't think of a better way to store and query a list type field without changing the schema of RegisteredModel. EntityAssociation is not really applicable for Prompts since it's managed by model registry store while EntityAssociation is in tracking store.

Makes sense, I think it is inevitable. Can we make sure this update is thread-safe and does not affect performance (I guess so cuz we file is async, but we should make sure no memory leak).

Sounds good, added a lock and moved the linkage logic so that it's called only when there's cache miss.

mlflow/tracking/_model_registry/fluent.py

B-Step62 · 2025-11-20T11:24:21Z

mlflow/tracking/client.py

+                    # Use LIKE to match the experiment ID and experiment ID is auto-incremented
+                    # integer. So add comma before and after the list of experiment IDs to
+                    # avoid false matches (e.g., "1" matches "10").
+                    exp_ids = f",{exp_ids},"


Makes sense, I think it is inevitable. Can we make sure this update is thread-safe and does not affect performance (I guess so cuz we file is async, but we should make sure no memory leak).

B-Step62 · 2025-11-20T11:25:00Z

/review check type hints and general coding standard

🚀 Review workflow started

✅ Review completed.

Review Output

Perfect! I've completed the PR review and added comments for all the issues I found. Let me summarize what I found:

PR Review Summary

I reviewed PR #18883 "Link Prompts and Experiments when prompts are loaded/registered" and found 3 issues that should be addressed:

Issues Found:

Security vulnerability (mlflow/store/model_registry/abstract_store.py:586): The LIKE pattern construction is vulnerable to SQL injection through wildcard injection. Since experiment IDs are auto-incremented integers, the code should validate that the experiment_id is numeric before using it in the query.
Documentation inconsistency (mlflow/tracking/_model_registry/fluent.py:667): The return type was changed from PagedList[Prompt] to list[Prompt], but the docstring still mentions pagination features like .token attribute. The documentation needs to be updated to match the actual return type.
Test reliability issue (tests/genai/prompts/test_prompts.py:1331): Multiple tests use time.sleep(1) to wait for async threading operations, which makes tests slower and potentially flaky. The codebase already has a join_thread_by_name_prefix() helper function that should be used instead for more reliable synchronization.

All comments have been posted to the PR with specific suggestions for fixes.

mlflow/store/model_registry/abstract_store.py

mlflow/tracking/_model_registry/fluent.py

tests/genai/prompts/test_prompts.py

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

…ent-link

… to experiments. Added a fixture to ensure linkage threads complete in tests. Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

Signed-off-by: TomuHirata <tomu.hirata@gmail.com>

B-Step62

LGTM!

tests/genai/prompts/test_prompts.py

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

TomeHirata · 2025-12-02T04:46:32Z

The failures are related to the conda issue.

…ow#18883) Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com> Signed-off-by: TomuHirata <tomu.hirata@gmail.com>

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com> Signed-off-by: TomuHirata <tomu.hirata@gmail.com>

TomeHirata added 4 commits November 18, 2025 13:51

Remove redundant import of PROMPT_EXPERIMENT_IDS_TAG_KEY in AbstractS…

1692b35

…tore class Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

Remove experiment_id filter syntax from documentation in multiple fil…

bb0c711

…es and update search_prompts function to return a list instead of PagedList. Enhance docstring for clarity on prompt search functionality. Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

github-actions bot added v3.6.1 area/prompts MLflow Prompt Registry and Optimization rn/feature Mention under Features in Changelogs. labels Nov 18, 2025

TomeHirata requested a review from Copilot November 18, 2025 08:05

Copilot started reviewing on behalf of TomeHirata November 18, 2025 08:06 View session

remove

fef1b7b

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

Copilot finished reviewing on behalf of TomeHirata November 18, 2025 08:08

Copilot AI reviewed Nov 18, 2025

View reviewed changes

TomeHirata commented Nov 18, 2025

View reviewed changes

TomeHirata requested a review from B-Step62 November 18, 2025 10:09

B-Step62 reviewed Nov 20, 2025

View reviewed changes

github-actions bot reviewed Nov 20, 2025

View reviewed changes

mlflow/store/model_registry/abstract_store.py Outdated Show resolved Hide resolved

github-actions bot reviewed Nov 20, 2025

View reviewed changes

mlflow/tracking/_model_registry/fluent.py Show resolved Hide resolved

github-actions bot reviewed Nov 20, 2025

View reviewed changes

tests/genai/prompts/test_prompts.py Outdated Show resolved Hide resolved

TomeHirata added 3 commits December 1, 2025 12:40

comments

71be495

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

Merge remote-tracking branch 'mlflow/master' into feat/prompt-experim…

37fedaa

…ent-link

Enhance error logging in prompt tagging and streamline prompt linking…

19010eb

… to experiments. Added a fixture to ensure linkage threads complete in tests. Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

TomeHirata requested a review from B-Step62 December 1, 2025 04:32

TomeHirata added the v3.7.0 label Dec 1, 2025

Merge branch 'master' into feat/prompt-experiment-link

d0dc7e1

Signed-off-by: TomuHirata <tomu.hirata@gmail.com>

B-Step62 approved these changes Dec 1, 2025

View reviewed changes

tests/genai/prompts/test_prompts.py Outdated Show resolved Hide resolved

lint

c1d4e23

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

TomeHirata enabled auto-merge December 2, 2025 01:16

TomeHirata disabled auto-merge December 2, 2025 04:46

TomeHirata merged commit b3f0e30 into mlflow:master Dec 2, 2025
43 of 47 checks passed

BenWilson2 pushed a commit to BenWilson2/mlflow that referenced this pull request Dec 4, 2025

Link Prompts and Experiments when prompts are loaded/registered (mlfl…

5325bce

…ow#18883) Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com> Signed-off-by: TomuHirata <tomu.hirata@gmail.com>

BenWilson2 pushed a commit that referenced this pull request Dec 4, 2025

Link Prompts and Experiments when prompts are loaded/registered (#18883)

45d90ce

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com> Signed-off-by: TomuHirata <tomu.hirata@gmail.com>

Conversation

TomeHirata commented Nov 18, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Install mlflow from this PR

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!

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!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Nov 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

TomeHirata Nov 18, 2025 • 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!

B-Step62 Nov 20, 2025

Choose a reason for hiding this comment

Uh oh!

TomeHirata Dec 1, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

B-Step62 Nov 20, 2025

Choose a reason for hiding this comment

Uh oh!

B-Step62 commented Nov 20, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Review Summary

Issues Found:

Uh oh!

Uh oh!

Uh oh!

Uh oh!

B-Step62 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

TomeHirata commented Dec 2, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

TomeHirata commented Nov 18, 2025 •

edited by github-actions bot

Loading

github-actions bot commented Nov 18, 2025 •

edited

Loading

TomeHirata Nov 18, 2025 •

edited

Loading

B-Step62 commented Nov 20, 2025 •

edited by github-actions bot

Loading