Add tracing integration to Gateway API endpoints by TomeHirata · Pull Request #20495 · mlflow/mlflow

TomeHirata · 2026-02-02T02:47:09Z

🥞 Stacked PR

Use this link to review incremental changes.

stack/gateway-trace-api [Files changed]
- stack/gateway/user-id [Files changed]
  - stack/gateway/passthrough-token [Files changed]

Related Issues/PRs

n/a

What changes are proposed in this pull request?

Title

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-02-02T03:10:15Z

🛠 DevTools 🛠

Install mlflow from this PR

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

github-actions · 2026-02-02T03:21:17Z

Documentation preview for cabbc7c is available at:

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

Copilot

Pull request overview

This PR adds first-class tracing / usage-tracking support for MLflow Gateway endpoints by persisting a tracing destination (experiment) on endpoints, instrumenting provider calls with spans (including streaming token usage where available), and exposing usage-tracking configuration in the Gateway UI.

Changes:

Add usage_tracking + experiment_id to Gateway endpoint persistence (DB schema, migrations, protos, entities, REST/SQL stores).
Instrument Gateway request handling and providers with MLflow tracing spans, including streamed token-usage extraction for select providers.
Extend UI and tests to configure usage tracking + select experiments and validate trace creation.

Reviewed changes

Copilot reviewed 47 out of 49 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
tests/tracking/test_rest_tracking.py	Adds coverage for auto-creating an experiment when `usage_tracking=True`.
tests/store/tracking/test_rest_store.py	Updates REST store request expectations to include `usage_tracking`.
tests/store/tracking/test_gateway_sql_store.py	Extends SQL store endpoint creation test with usage-tracking fields.
tests/server/test_gateway_api.py	Updates gateway API tests for traced providers, streaming, and trace creation.
tests/resources/db/latest_schema.sql	Updates latest test DB schema with new endpoint columns.
tests/gateway/schemas/test_completions.py	Validates streaming completions schema with optional `usage`.
tests/gateway/schemas/test_chat.py	Validates streaming chat schema with optional `usage`.
tests/gateway/providers/test_tracing.py	Adds unit tests for provider tracing wrapper behavior.
tests/gateway/providers/test_togetherai.py	Updates TogetherAI streaming expectations to include `usage: None` / usage final chunk.
tests/gateway/providers/test_openai.py	Updates OpenAI streaming tests for `stream_options.include_usage` + optional usage field.
tests/gateway/providers/test_gemini.py	Updates Gemini streaming tests to include optional `usage`.
tests/gateway/providers/test_cohere.py	Updates Cohere streaming tests to include usage objects.
tests/gateway/providers/test_anthropic.py	Updates Anthropic streaming tests + validates usage extraction changes.
tests/db/schemas/sqlite.sql	Adds `experiment_id` and `usage_tracking` to SQLite test schema.
tests/db/schemas/postgresql.sql	Adds `experiment_id` and `usage_tracking` to Postgres test schema.
tests/db/schemas/mysql.sql	Adds `experiment_id` and `usage_tracking` to MySQL test schema.
tests/db/schemas/mssql.sql	Adds `experiment_id` and `usage_tracking` to MSSQL test schema.
mlflow/types/chat.py	Adds optional `usage` field to chat completion chunk type.
mlflow/tracing/constant.py	Adds gateway-related trace metadata keys + provider/model span attributes.
mlflow/store/tracking/gateway/sqlalchemy_mixin.py	Persists `experiment_id` + `usage_tracking` for endpoint create/update.
mlflow/store/tracking/gateway/rest_mixin.py	Extends REST store endpoint create/update payloads with usage-tracking fields.
mlflow/store/tracking/gateway/entities.py	Adds `experiment_id` to resolved endpoint config entity.
mlflow/store/tracking/gateway/config_resolver.py	Propagates `experiment_id` into runtime endpoint config.
mlflow/store/tracking/gateway/abstract_mixin.py	Updates abstract store contract/docs for usage tracking + experiment handling.
mlflow/store/tracking/dbmodels/models.py	Adds new endpoint columns to SQLAlchemy DB model and entity conversion.
mlflow/store/db_migrations/versions/d0e1f2a3b4c5_add_experiment_id_to_endpoints.py	Adds Alembic migration for new endpoint columns.
mlflow/server/js/src/lang/default/en.json	Adds UI strings for usage tracking + experiment selection + “View traces”.
mlflow/server/js/src/gateway/types.ts	Extends Gateway endpoint/request types with usage-tracking fields.
mlflow/server/js/src/gateway/pages/EndpointPage.tsx	Passes experiment id into edit form renderer to surface trace link.
mlflow/server/js/src/gateway/hooks/useExperimentsForSelect.ts	Adds query hook to fetch experiments for selection.
mlflow/server/js/src/gateway/hooks/useCreateEndpointForm.ts	Adds usage tracking + experiment selection to create-endpoint submit payload.
mlflow/server/js/src/gateway/components/endpoint-form/EndpointFormRenderer.tsx	Adds create-time usage tracking toggle + experiment selector UI.
mlflow/server/js/src/gateway/components/edit-endpoint/EditEndpointFormRenderer.tsx	Adds “View traces” link when endpoint has an experiment id.
mlflow/server/js/src/gateway/components/create-endpoint/ExperimentSelect.tsx	Adds experiment select component for the create form.
mlflow/server/handlers.py	Adds server-side experiment auto-creation logic for usage-tracked endpoints.
mlflow/server/gateway_api.py	Adds gateway span creation + traced streaming response wrapper; wraps providers for tracing.
mlflow/protos/service_pb2.pyi	Updates Python proto stubs for new endpoint fields.
mlflow/protos/service.proto	Adds `experiment_id` + `usage_tracking` to gateway endpoint create/update protos.
mlflow/gateway/utils.py	Updates SSE serialization to use `model_dump_json()` for pydantic v2.
mlflow/gateway/schemas/completions.py	Adds optional `usage` to streaming completions response schema.
mlflow/gateway/providers/tracing.py	Introduces `TracingProviderWrapper` for provider method span instrumentation.
mlflow/gateway/providers/openai.py	Adds streamed usage extraction + injects `stream_options.include_usage`.
mlflow/gateway/providers/litellm.py	Adds `get_provider_name()` override for more accurate tracing labels.
mlflow/gateway/providers/gemini.py	Adds streamed usage extraction from `usageMetadata`.
mlflow/gateway/providers/base.py	Adds `get_provider_name()` default to support tracing/metrics naming.
mlflow/gateway/providers/anthropic.py	Adds streamed usage aggregation across events and exposes it on chunks.
mlflow/entities/gateway_endpoint.py	Extends GatewayEndpoint entity (proto conversion) with new fields.

mlflow/server/gateway_api.py

mlflow/gateway/providers/tracing.py

mlflow/store/tracking/gateway/sqlalchemy_mixin.py

mlflow/server/handlers.py

Copilot

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

mlflow/server/gateway_api.py

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

mlflow/gateway/tracing_utils.py

harupy · 2026-02-05T02:11:31Z

/review

🚀 Review workflow started

✅ Review completed.

Review Output

The code looks consistent with the implementation. Let me now finalize my review. After careful examination of all the changed lines in this PR, I have reviewed:

mlflow/gateway/tracing_utils.py - New utility module
mlflow/server/gateway_api.py - Integration changes
tests/server/test_gateway_api.py - Test updates
tests/telemetry/test_tracked_events.py - Test updates

After thorough review of the changed lines only, here's my assessment:

No issues found

The PR changes are well-structured:

tracing_utils.py: The traced_gateway_call function correctly handles the case when tracing is disabled (no experiment_id) by returning the original function unchanged, and properly wraps it with mlflow.trace when enabled.
gateway_api.py: The changes consistently update all endpoint handlers to:
- Unpack the new tuple return from _create_provider_from_endpoint_name
- Use traced_gateway_call to wrap provider methods
- Replace the removed make_streaming_response with inline StreamingResponse using to_sse_chunk
Test updates: Tests properly mock the new tuple return value (mock_provider, mock_endpoint_config) and update streaming test assertions to work with the new implementation.
The import changes: make_streaming_response is removed and replaced with to_sse_chunk and inline StreamingResponse, which is cleaner.
The change from body.get("stream") to body.get("stream", False) is a minor improvement that makes the default explicit.

The code follows consistent patterns, has proper type hints, and the tests comprehensively cover the new tracing functionality including both synchronous and streaming cases.

No issues found

tests/server/test_gateway_api.py

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

serena-ruan · 2026-02-05T06:42:43Z

mlflow/server/gateway_api.py

    endpoint_type: EndpointType,
    enable_tracing: bool = True,
-) -> BaseProvider:
+) -> tuple[BaseProvider, GatewayEndpointConfig]:


#20495 (comment)

This is surprising that we have EndpointConfig and GatewayEndpointConfig, it's confusing. I feel like we should rename EndpointConfig to ProviderConfig but not related to this PR.

Could we rename this function though, and probably use a NamedTuple for return value to be more explicit

Yeah, I think we'll consolidate them once we deprecate or integrate the legacy gateway workflow

serena-ruan · 2026-02-05T06:45:53Z

mlflow/gateway/tracing_utils.py

+from mlflow.store.tracking.gateway.entities import GatewayEndpointConfig
+
+
+def traced_gateway_call(


nit for this function name, since it may or may not be traced depending on the config, maybe something like apply_gateway_tracing_config

maybe_traced_gateway_call?

serena-ruan

Overall LGTM! Left some nits :)

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

This was referenced Feb 2, 2026

link gateway and experiment #20356

Merged

Add trace ingestion for Gateway endpoints #20358

Merged

Add usage section in endpoint page #20357

Merged

TomeHirata force-pushed the stack/gateway-trace-api branch from c85cd23 to f62f74c Compare February 2, 2026 02:59

TomeHirata marked this pull request as ready for review February 2, 2026 03:10

Copilot AI review requested due to automatic review settings February 2, 2026 03:10

TomeHirata requested a review from serena-ruan February 2, 2026 03:10

Copilot started reviewing on behalf of TomeHirata February 2, 2026 03:10 View session

TomeHirata force-pushed the stack/gateway-trace-api branch from f62f74c to 8f4c23a Compare February 2, 2026 03:11

Copilot AI reviewed Feb 2, 2026

View reviewed changes

TomeHirata force-pushed the stack/gateway-trace-api branch 2 times, most recently from 2d9066d to acf47cd Compare February 2, 2026 05:42

TomeHirata mentioned this pull request Feb 2, 2026

Add usage tracking toggle and experiment selector to endpoint creation form #20497

Merged

29 tasks

TomeHirata commented Feb 2, 2026

View reviewed changes

mlflow/server/gateway_api.py Outdated Show resolved Hide resolved

TomeHirata commented Feb 2, 2026

View reviewed changes

mlflow/server/gateway_api.py Outdated Show resolved Hide resolved

TomeHirata force-pushed the stack/gateway-trace-api branch 12 times, most recently from 0a1d073 to b18ba2f Compare February 3, 2026 07:21

TomeHirata force-pushed the stack/gateway-trace-api branch from de9b744 to 8368e29 Compare February 4, 2026 05:17

comments

4c2b9fe

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

TomeHirata requested a review from serena-ruan February 4, 2026 06:34

remove unused comment

b0833d2

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

TomeHirata mentioned this pull request Feb 4, 2026

Add user attributes to gateway span #20559

Merged

31 tasks

TomeHirata added 2 commits February 4, 2026 22:44

test

8586580

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

test

38f76a3

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

TomeHirata mentioned this pull request Feb 5, 2026

Support token usage tracking for passthrough #20578

Merged

31 tasks

TomeHirata added the team-review Trigger a team review request label Feb 5, 2026

github-actions bot requested review from daniellok-db and harupy February 5, 2026 01:16

test

a7104da

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

github-actions bot added the rn/none List under Small Changes in Changelogs. label Feb 5, 2026

harupy reviewed Feb 5, 2026

View reviewed changes

mlflow/gateway/tracing_utils.py Outdated Show resolved Hide resolved

harupy reviewed Feb 5, 2026

View reviewed changes

tests/server/test_gateway_api.py Outdated Show resolved Hide resolved

github-actions bot assigned harupy Feb 5, 2026

comments

a693bb9

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

TomeHirata requested a review from harupy February 5, 2026 04:55

harupy approved these changes Feb 5, 2026

View reviewed changes

serena-ruan reviewed Feb 5, 2026

View reviewed changes

serena-ruan approved these changes Feb 5, 2026

View reviewed changes

TomeHirata added 2 commits February 5, 2026 16:19

comments

a3b1ac7

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

test

cabbc7c

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

TomeHirata enabled auto-merge February 5, 2026 07:32

TomeHirata added this pull request to the merge queue Feb 5, 2026

Merged via the queue into mlflow:master with commit 4eb654b Feb 5, 2026
52 checks passed

TomeHirata deleted the stack/gateway-trace-api branch February 5, 2026 08:06

		from mlflow.store.tracking.gateway.entities import GatewayEndpointConfig


		def traced_gateway_call(

Conversation

TomeHirata commented Feb 2, 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?

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 Feb 2, 2026

Install mlflow from this PR

Uh oh!

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

Uh oh!

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!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

Uh oh!

serena-ruan Feb 5, 2026

Choose a reason for hiding this comment

Uh oh!

TomeHirata Feb 5, 2026

Choose a reason for hiding this comment

Uh oh!

serena-ruan Feb 5, 2026

Choose a reason for hiding this comment

Uh oh!

TomeHirata Feb 5, 2026

Choose a reason for hiding this comment

Uh oh!

serena-ruan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

TomeHirata commented Feb 2, 2026 •

edited

Loading

github-actions bot commented Feb 2, 2026 •

edited

Loading

harupy commented Feb 5, 2026 •

edited by github-actions bot

Loading