Support mlflow.genai.to_predict_fn for app invocation endpoints by jennsun · Pull Request #19779 · mlflow/mlflow

jennsun · 2026-01-06T22:52:52Z

🛠 DevTools 🛠

Install mlflow from this PR

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

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

This PR updates genai.to_predict_fn (https://mlflow.org/docs/latest/api_reference/python_api/mlflow.genai.html#mlflow.genai.to_predict_fn) to support calls to app invocation endpoints

Where a user can create a predict function by using

predict_fn = to_predict_fn(f"apps:/{APP_NAME}")

adding an additional schema to the existing group (used to only support model serving endpoints with schema endpoints)

Under the hood, an oauth token for the databricks workspace will be used to send a post request to the app's invocation endpoint to be passed in as the prediction function for mlflow e2e eval flow

How is this PR tested?

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

Tested notebook with mlflow version of this PR [NOTE: app invocations requires a fix for notebook oauth tokens]:
https://eng-ml-agent-platform.staging.cloud.databricks.com/editor/notebooks/2722550516530718?o=2850744067564480

Tested locally:

With mlflow evaluation framework, verifying traces are visible in experiments:

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
[ X ] 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)

mlflow/utils/databricks_utils.py

github-actions · 2026-01-07T02:15:44Z

Documentation preview for 6e11bbc is available at:

https://pr-19779--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.

github-actions · 2026-01-07T17:19:08Z

@jennsun 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.

bbqiu · 2026-01-08T01:15:53Z

mlflow/genai/evaluation/base.py


+    schema, path = _parse_model_uri(endpoint_uri)
+
+    if schema == "apps":


nit: can we separate this function into 2 helper functions?

+1, let's use case match for schema and route to different functions

mlflow/utils/databricks_utils.py

bbqiu · 2026-01-08T01:31:44Z

tests/genai/evaluate/test_to_predict_fn.py

        mock_get_experiment_id.assert_called_once()


+# ========== Databricks Apps Tests ==========


for the manual test, can we do a whole e2e flow for eval?

e2e notebook example: https://eng-ml-agent-platform.staging.cloud.databricks.com/editor/notebooks/2722550516530718?o=2850744067564480#command/5285377916446537

experiment results with traces: https://eng-ml-agent-platform.staging.cloud.databricks.com/ml/experiments/2722550516530718/evaluation-runs?selectedRunUuid=4fac0bf652c044b598e1505adc26f77c&o=2850744067564480&selectedColumns=trace_id%2Cexecution_duration%2Cstate

bbqiu · 2026-01-08T01:32:44Z

mlflow/genai/evaluation/base.py

+
+    if schema == "apps":
+        try:
+            config = get_databricks_workspace_client_config("databricks", scopes=["all-apis"])


nit: can we leave a comment here describing why this will be okay?

bbqiu

great work! left a few comments

you might have to rebase and sign all of your commits btw

in the future you can avoid by running git commit -sm .. to sign each commit

Signed-off-by: Jenny <jenny.sun@databricks.com>

jennsun · 2026-01-09T18:06:46Z

Note: upon further testing with notebook oauth it looks like the notebook oauth tokens are failing to query the invocations endpoint for apps in notebook environment due to oauth2proxy failures:

[2026/01/09 01:22:38] [databricks_apps.go:536] ActorID not found in session user, will attempt to get from current user API
[2026/01/09 01:22:38] [jwt_session.go:52] Error retrieving session from token in Authorization header: [unable to verify bearer token, could not check roles: error getting current user: {"error_description":"Invalid certificate confirmation","error":"access_denied"} [ReqId: 12713ec0-415b-4568-9564-a1c6dfbe4e74]]
[2026/01/09 01:22:38] [logger.go:777] [req=5d1377e7-cd0e-4a04-8c36-77fc1ed94c5e] No valid authentication in request. Initiating login.

This happens because the notebook oauth tokens contain a CNF/X.509 Fingerprint (SHA-256) claim which binds the token to a specific client certificate (intended for notebook usage only). oauth2proxy is not presenting the right certificate when making API calls, so when it tries to use this token without presenting the certificate, the request is rejected and the endpoint redirects to sign-in page.

mlflow/genai/evaluation/base.py

serena-ruan · 2026-01-12T09:58:25Z

mlflow/genai/evaluation/base.py


+    schema, path = _parse_model_uri(endpoint_uri)
+
+    if schema == "apps":


+1, let's use case match for schema and route to different functions

mlflow/utils/databricks_utils.py

Signed-off-by: Jenny <jenny.sun@databricks.com>

mlflow/genai/evaluation/base.py

Signed-off-by: Jenny <jenny.sun@databricks.com>

mlflow/genai/evaluation/base.py

Co-authored-by: Serena Ruan <82044803+serena-ruan@users.noreply.github.com> Signed-off-by: Jenny <jennysunnjm@gmail.com>

serena-ruan · 2026-01-14T02:48:19Z

mlflow/utils/databricks_utils.py

+    except TypeError as e:
+        if "scopes" in str(e):
+            raise MlflowException.invalid_parameter_value(
+                "The 'scopes' parameter requires databricks-sdk>=0.74.0. "
+                "Please upgrade with: pip install --upgrade databricks-sdk",
+            ) from e
+        raise


Sorry if I didn't make it clear in previous comment, could we check databricks sdk version instead of capturing this error message, which is not reliable? It may capture other errors with scopes but not related to version.
Let's create a function like:

def databricks_workspace_client_supports_scopes: if Version(databricks.sdk.xxx) < Version("0.74.0"): raise ...

Then we can reuse this check here and line382 in models/evaluation/base.py

tests/genai/evaluate/test_to_predict_fn.py

serena-ruan

Overall LGTM! BTW we pushed RC to tomorrow, so you can address the comment before then :)

Signed-off-by: Jenny <jenny.sun@databricks.com>

serena-ruan · 2026-01-14T04:19:44Z

mlflow/utils/databricks_utils.py

+    Raises:
+        MlflowException: If databricks-sdk version is < 0.74.0
+    """
+    from packaging.version import Version


Let's move this to top

mlflow/utils/databricks_utils.py

serena-ruan · 2026-01-14T04:22:03Z

mlflow/utils/databricks_utils.py

    from databricks.sdk import WorkspaceClient

+    if scopes is not None:
+        check_databricks_sdk_supports_scopes()


Let's use kwargs = {"scopes": scopes} here, and pass below to avoid breaking old databricks-sdk

Signed-off-by: Jenny <jenny.sun@databricks.com>

bbqiu · 2026-01-14T22:36:42Z

mlflow/genai/evaluation/base.py

+
+    # Append /invocations to endpoint
+    return DatabricksAppConfig(
+        app_url=f"{app.url}/invocations",


nit: can we change this to be app_invocation_url?

bbqiu · 2026-01-14T22:38:39Z

mlflow/utils/databricks_utils.py

+
+def check_databricks_sdk_supports_scopes():
+    """
+    Check if the installed databricks-sdk version supports the 'scopes' parameter.


nit: specify param for workspaceclient

jennsun · 2026-01-14T22:40:06Z

mlflow/genai/evaluation/base.py

+    Returns:
+        A predict function that invokes the endpoint
+    """
+    from mlflow.deployments import get_deploy_client


nit: move imports up

bbqiu

LGTM

Signed-off-by: Jenny <jenny.sun@databricks.com>

…ow#19779) Signed-off-by: Jenny <jenny.sun@databricks.com> Signed-off-by: Jenny <jennysunnjm@gmail.com> Co-authored-by: Serena Ruan <82044803+serena-ruan@users.noreply.github.com>

jennsun changed the title ~~[draft] apps to_predict_fn~~ Support mlflow.genai.to_predict_fn for app invocation endpoints Jan 7, 2026

jennsun commented Jan 7, 2026

View reviewed changes

mlflow/utils/databricks_utils.py Outdated Show resolved Hide resolved

jennsun marked this pull request as ready for review January 7, 2026 02:08

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

jennsun marked this pull request as draft January 7, 2026 02:08

jennsun marked this pull request as ready for review January 7, 2026 17:18

jennsun force-pushed the ML-57538-app-to-predict-fn branch from 7297367 to a108b51 Compare January 7, 2026 17:22

bbqiu self-requested a review January 7, 2026 19:15

bbqiu reviewed Jan 8, 2026

View reviewed changes

mlflow/utils/databricks_utils.py Outdated Show resolved Hide resolved

bbqiu reviewed Jan 8, 2026

View reviewed changes

mlflow/utils/databricks_utils.py Outdated Show resolved Hide resolved

bbqiu reviewed Jan 8, 2026

View reviewed changes

jennsun force-pushed the ML-57538-app-to-predict-fn branch from a108b51 to 9c9dfb2 Compare January 9, 2026 00:37

add apps to predict fn to call invocations endpoint

a2a3319

Signed-off-by: Jenny <jenny.sun@databricks.com>

jennsun force-pushed the ML-57538-app-to-predict-fn branch from 9c9dfb2 to a2a3319 Compare January 9, 2026 00:38

jennsun requested a review from serena-ruan January 11, 2026 16:28

serena-ruan reviewed Jan 12, 2026

View reviewed changes

github-actions bot assigned serena-ruan Jan 12, 2026

mlflow pr review

1581a3d

Signed-off-by: Jenny <jenny.sun@databricks.com>

jennsun requested a review from serena-ruan January 12, 2026 23:13

pass scopes into tests/revert mse docstring

3b1f1a3

Signed-off-by: Jenny <jenny.sun@databricks.com>

serena-ruan reviewed Jan 13, 2026

View reviewed changes

mlflow/genai/evaluation/base.py Outdated Show resolved Hide resolved

serena-ruan reviewed Jan 13, 2026

View reviewed changes

mlflow/genai/evaluation/base.py Outdated Show resolved Hide resolved

serena-ruan reviewed Jan 13, 2026

View reviewed changes

mlflow/genai/evaluation/base.py Outdated Show resolved Hide resolved

serena-ruan reviewed Jan 13, 2026

View reviewed changes

mlflow/genai/evaluation/base.py Outdated Show resolved Hide resolved

jennsun added 2 commits January 13, 2026 08:11

pr review named tuple sdk updates

d40e29e

Signed-off-by: Jenny <jenny.sun@databricks.com>

lint

e50a94b

Signed-off-by: Jenny <jenny.sun@databricks.com>

jennsun requested a review from serena-ruan January 13, 2026 18:50

serena-ruan added the v3.9.0 label Jan 14, 2026

serena-ruan reviewed Jan 14, 2026

View reviewed changes

mlflow/genai/evaluation/base.py Outdated Show resolved Hide resolved

Update mlflow/genai/evaluation/base.py

f7add3d

Co-authored-by: Serena Ruan <82044803+serena-ruan@users.noreply.github.com> Signed-off-by: Jenny <jennysunnjm@gmail.com>

serena-ruan reviewed Jan 14, 2026

View reviewed changes

tests/genai/evaluate/test_to_predict_fn.py Outdated Show resolved Hide resolved

serena-ruan approved these changes Jan 14, 2026

View reviewed changes

jennsun added 3 commits January 13, 2026 19:10

pr review - check sdk version, directly test traces

888aef2

Signed-off-by: Jenny <jenny.sun@databricks.com>

sdk test check databricks-sdk

6db8065

Signed-off-by: Jenny <jenny.sun@databricks.com>

lint

26c03fe

Signed-off-by: Jenny <jenny.sun@databricks.com>

serena-ruan reviewed Jan 14, 2026

View reviewed changes

mlflow/utils/databricks_utils.py Outdated Show resolved Hide resolved

serena-ruan reviewed Jan 14, 2026

View reviewed changes

jennsun added 3 commits January 13, 2026 20:34

pass scopes in as kwargs

e189992

Signed-off-by: Jenny <jenny.sun@databricks.com>

fix tests

a7101f9

Signed-off-by: Jenny <jenny.sun@databricks.com>

fix tests

4b5a008

Signed-off-by: Jenny <jenny.sun@databricks.com>

bbqiu reviewed Jan 14, 2026

View reviewed changes

jennsun commented Jan 14, 2026

View reviewed changes

bbqiu approved these changes Jan 14, 2026

View reviewed changes

jennsun added 3 commits January 14, 2026 15:08

nits

c04808b

Signed-off-by: Jenny <jenny.sun@databricks.com>

test fixes

09bc7e5

Signed-off-by: Jenny <jenny.sun@databricks.com>

Merge branch 'master' into ML-57538-app-to-predict-fn

6e11bbc

jennsun added this pull request to the merge queue Jan 15, 2026

Merged via the queue into mlflow:master with commit 7d256d2 Jan 15, 2026
48 of 50 checks passed

jennsun deleted the ML-57538-app-to-predict-fn branch January 15, 2026 02:20


		schema, path = _parse_model_uri(endpoint_uri)

		if schema == "apps":

		mock_get_experiment_id.assert_called_once()


		# ========== Databricks Apps Tests ==========

Conversation

jennsun commented Jan 6, 2026 • edited 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!

Uh oh!

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

Uh oh!

Uh oh!

github-actions bot commented Jan 7, 2026

⚠ DCO check

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

bbqiu Jan 8, 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!

bbqiu left a comment

Choose a reason for hiding this comment

Uh oh!

jennsun commented Jan 9, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

serena-ruan left a comment

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!

Choose a reason for hiding this comment

Uh oh!

bbqiu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

jennsun commented Jan 6, 2026 •

edited

Loading

github-actions bot commented Jan 7, 2026 •

edited

Loading

bbqiu Jan 8, 2026 •

edited

Loading