Remove databricks serving endpoint adapter and redirect databricks managed models to litellm adapter by xsh310 · Pull Request #20037 · mlflow/mlflow

xsh310 · 2026-01-15T21:13:57Z

🥞 Stacked PR

Use this link to review incremental changes.

stack/remove-databricks_serving_endpoint_adapter-and-direct-databricks-managed-models-to-litellm-adapter [Files changed]

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Remove the databricks serving endpoint adapter to reduce the number of code paths that we have to manage. Databricks managed models will be redirected to litellm adapter instead.

How is this PR tested?

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

Manual Testing

Tested with this notebook: https://e2-dogfood.staging.cloud.databricks.com/editor/notebooks/3782812667401322?o=6051921418418893

Verified that the following is working:

LLM judge invocation: Databrick-manged OpenAI models, Gemini models, Claude models
Agent-as-judge: Databrick-managed OpenAI, Gemini models
create_deepeval_model and create_ragas_model still works
Conversation Simulation with databricks model

Verified that the following is not working:
Agent-as-judge: Databrick-managed Claude model (Filed a ticket to follow up)

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-15T22:11:07Z

🛠 DevTools 🛠

Install mlflow from this PR

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

github-actions · 2026-01-15T22:19:07Z

Documentation preview for e671f7b is available at:

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

smoorjani

left a few comments about telemetry and tests, but mostly LGTM. Thanks for the fix and extensive testing!

smoorjani · 2026-01-16T00:32:01Z

tests/genai/judges/adapters/test_databricks_managed_adapter.py

-            {"databricks.agents.telemetry": mock_telemetry_module},
-        ),
-    ):
-        _record_judge_model_usage_success_databricks_telemetry(


I don't think we should be getting rid of this, right? Two points here:

It seems like this should also be in the managed adapter (doesn't seem like it is right now)

We should still have this in the case of a databricks:/... endpoint

Makes sense, mirrored the telemetry related ones in test_litellm_adapter.py

smoorjani · 2026-01-16T21:29:08Z

tests/genai/judges/utils/test_invocation_utils.py

-    ],
-)
-@pytest.mark.parametrize("with_trace", [False, True])
-def test_invoke_judge_model_databricks_success_not_in_databricks(


can we ensure we still have coverage for databricks model URIs via invoke_judge_model? similar to the note above - we should ensure that the appropriate telemetry (and any other databricks-specific behaviors) are being triggered.

smoorjani · 2026-01-16T21:29:43Z

tests/genai/scorers/deepeval/test_utils.py

-    assert model.__class__.__name__ == "DatabricksServingEndpointDeepEvalLLM"
-    assert model.get_model_name() == "databricks:/my-endpoint"
+    assert model.__class__.__name__ == "LiteLLMModel"
+    assert "my-endpoint" in model.get_model_name()


can we assert on exact format instead of using in? We may also want to check provider if this is split up

smoorjani · 2026-01-16T21:29:56Z

tests/genai/scorers/phoenix/test_models.py


 from mlflow.exceptions import MlflowException
+
+# Import phoenix.evals - skip tests if not available or incompatible version


we should not skip tests

hmm good catch, claude must have changed this after it was unable to import...

xsh310 · 2026-01-17T02:22:19Z

mlflow/genai/judges/adapters/litellm_adapter.py

 _suppress_litellm_nonfatal_errors = _SuppressLiteLLMNonfatalErrors()


+def _record_judge_model_usage_success_databricks_telemetry(


Aligned with @smoorjani offline that we will keep the telemetry for databricks-served models for now in litellm adapter before we reassess whether it is still needed.

smoorjani

LGTM, left two minor comments to address before merging. Thanks!

smoorjani · 2026-01-17T15:57:34Z

mlflow/genai/judges/adapters/litellm_adapter.py

+                usage = getattr(response, "usage", None)
+                prompt_tokens = getattr(usage, "prompt_tokens", None) if usage else None
+                completion_tokens = getattr(usage, "completion_tokens", None) if usage else None
+                return message.content, total_cost, request_id, prompt_tokens, completion_tokens


can we use a dataclass for this instead of returning a tuple?

smoorjani · 2026-01-17T16:06:25Z

tests/genai/judges/adapters/test_litellm_adapter.py

+        mock_job_run_id.assert_called_once()
+
+
+def test_record_success_telemetry_without_databricks_agents():


is it possible to assert that nothing was recorded? can we also add this test for failure telemetry?

dbczumar

LGTM! Can we get approval from @B-Step62 as well?

B-Step62

LGTM!

B-Step62 · 2026-01-21T09:03:31Z

mlflow/genai/judges/adapters/litellm_adapter.py

 _suppress_litellm_nonfatal_errors = _SuppressLiteLLMNonfatalErrors()


+def _record_judge_model_usage_success_databricks_telemetry(


Can we move these telemetry methods out to a separate util?

…naged models to litellm adapter Signed-off-by: Xiang Shen <xshen.shc@gmail.com>

Signed-off-by: Xiang Shen <xshen.shc@gmail.com>

…y tests Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com> Signed-off-by: Xiang Shen <xshen.shc@gmail.com>

…lpers out of litellm adapter into its own util file Signed-off-by: Xiang Shen <xshen.shc@gmail.com>

xsh310 · 2026-01-21T20:51:19Z

Updated PR to move telemetry methods to its own util file

AveshCSingh · 2026-01-21T23:30:13Z

I'm going to force merge this PR since the CI failure of test_projects_cli.py is unrelated to this change.

[ML-61124]: update translation labels [ML-61124]: add feature flag Remove databricks serving endpoint adapter and redirect databricks managed models to litellm adapter (mlflow#20037) Signed-off-by: Xiang Shen <xshen.shc@gmail.com> Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>

[ML-61124]: update translation labels [ML-61124]: add feature flag Remove databricks serving endpoint adapter and redirect databricks managed models to litellm adapter (mlflow#20037) Signed-off-by: Xiang Shen <xshen.shc@gmail.com> Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com> [ML-61125]: set shouldEnableWorkflowBasedNavigation to false

[ML-61124]: update translation labels [ML-61124]: add feature flag Remove databricks serving endpoint adapter and redirect databricks managed models to litellm adapter (mlflow#20037) Signed-off-by: Xiang Shen <xshen.shc@gmail.com> Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com> [ML-61125]: set shouldEnableWorkflowBasedNavigation to false [ML-61124]: sidebar width

[ML-61124]: update translation labels [ML-61124]: add feature flag Remove databricks serving endpoint adapter and redirect databricks managed models to litellm adapter (mlflow#20037) Signed-off-by: Xiang Shen <xshen.shc@gmail.com> Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com> [ML-61125]: set shouldEnableWorkflowBasedNavigation to false [ML-61124]: sidebar width [ML-61125]: resolve conflict

[ML-61124]: update translation labels [ML-61124]: add feature flag Remove databricks serving endpoint adapter and redirect databricks managed models to litellm adapter (mlflow#20037) Signed-off-by: Xiang Shen <xshen.shc@gmail.com> Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com> [ML-61125]: set shouldEnableWorkflowBasedNavigation to false [ML-61124]: sidebar width

[ML-61124]: update translation labels [ML-61124]: add feature flag Remove databricks serving endpoint adapter and redirect databricks managed models to litellm adapter (mlflow#20037) Signed-off-by: Xiang Shen <xshen.shc@gmail.com> Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com> [ML-61125]: set shouldEnableWorkflowBasedNavigation to false [ML-61124]: sidebar width [ML-61125]: resolve conflict

[ML-61124]: update translation labels [ML-61124]: add feature flag Remove databricks serving endpoint adapter and redirect databricks managed models to litellm adapter (mlflow#20037) Signed-off-by: Xiang Shen <xshen.shc@gmail.com> Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com> [ML-61125]: set shouldEnableWorkflowBasedNavigation to false [ML-61124]: sidebar width

…naged models to litellm adapter (mlflow#20037) Signed-off-by: Xiang Shen <xshen.shc@gmail.com> Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>

…naged models to litellm adapter (#20037) Signed-off-by: Xiang Shen <xshen.shc@gmail.com> Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>

xsh310 changed the title ~~remove databricks serving endpoint adapter and redirect databricks managed models to litellm adapter~~ Remove databricks serving endpoint adapter and redirect databricks managed models to litellm adapter Jan 15, 2026

xsh310 marked this pull request as ready for review January 15, 2026 22:10

xsh310 requested review from AveshCSingh, B-Step62, dbczumar and smoorjani January 15, 2026 22:10

github-actions bot added area/evaluation MLflow Evaluation rn/none List under Small Changes in Changelogs. labels Jan 15, 2026

xsh310 force-pushed the stack/remove-databricks_serving_endpoint_adapter-and-direct-databricks-managed-models-to-litellm-adapter branch from 1b384a8 to 062003c Compare January 15, 2026 22:28

github-actions bot added the v3.8.2 label Jan 15, 2026

xsh310 force-pushed the stack/remove-databricks_serving_endpoint_adapter-and-direct-databricks-managed-models-to-litellm-adapter branch from 062003c to a5cba14 Compare January 15, 2026 22:50

smoorjani requested changes Jan 16, 2026

View reviewed changes

github-actions bot assigned smoorjani Jan 17, 2026

xsh310 force-pushed the stack/remove-databricks_serving_endpoint_adapter-and-direct-databricks-managed-models-to-litellm-adapter branch 2 times, most recently from ffea371 to c417406 Compare January 17, 2026 02:17

xsh310 commented Jan 17, 2026

View reviewed changes

smoorjani approved these changes Jan 17, 2026

View reviewed changes

dbczumar added the v3.9.0 label Jan 20, 2026

dbczumar approved these changes Jan 20, 2026

View reviewed changes

github-actions bot assigned dbczumar Jan 21, 2026

B-Step62 approved these changes Jan 21, 2026

View reviewed changes

xsh310 and others added 4 commits January 21, 2026 12:37

remove databricks serving endpoint adapter and redirect databricks ma…

9a53379

…naged models to litellm adapter Signed-off-by: Xiang Shen <xshen.shc@gmail.com>

address comments on telemetry and tests

c3aa851

Signed-off-by: Xiang Shen <xshen.shc@gmail.com>

Replace tuple with InvokeLiteLLMOutput dataclass and improve telemetr…

e682385

…y tests Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com> Signed-off-by: Xiang Shen <xshen.shc@gmail.com>

Move record judge model usage success/failure databricks telemetry he…

e671f7b

…lpers out of litellm adapter into its own util file Signed-off-by: Xiang Shen <xshen.shc@gmail.com>

xsh310 force-pushed the stack/remove-databricks_serving_endpoint_adapter-and-direct-databricks-managed-models-to-litellm-adapter branch from ceaafbf to e671f7b Compare January 21, 2026 20:49

xsh310 enabled auto-merge January 21, 2026 20:56

github-actions bot assigned B-Step62 Jan 21, 2026

xsh310 disabled auto-merge January 21, 2026 23:17

AveshCSingh merged commit 6ed06c4 into mlflow:master Jan 21, 2026
82 of 86 checks passed


		from mlflow.exceptions import MlflowException

		# Import phoenix.evals - skip tests if not available or incompatible version

		_suppress_litellm_nonfatal_errors = _SuppressLiteLLMNonfatalErrors()


		def _record_judge_model_usage_success_databricks_telemetry(

		mock_job_run_id.assert_called_once()


		def test_record_success_telemetry_without_databricks_agents():

Conversation

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

Uh oh!

🥞 Stacked PR

Related Issues/PRs

What changes are proposed in this pull request?

How is this PR tested?

Manual Testing

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

Install mlflow from this PR

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!

smoorjani left a comment

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!

smoorjani left a comment

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!

dbczumar left a comment

Choose a reason for hiding this comment

Uh oh!

B-Step62 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xsh310 commented Jan 21, 2026

Uh oh!

AveshCSingh commented Jan 21, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

xsh310 commented Jan 15, 2026 •

edited

Loading

github-actions bot commented Jan 15, 2026 •

edited

Loading