Distributed tracing for gateway endpoints by TomeHirata · Pull Request #20864 · mlflow/mlflow

TomeHirata · 2026-02-16T10:53:50Z

Related Issues/PRs

n/a

What changes are proposed in this pull request?

When the gateway endpoint is called by a traced agent, we want to show gateway execution in the trace as well. To achieve both gateway endpoint observability and the distributed tracing between client and gateway, we will duplicate the gateway spans but without input/output payloads to minimize the additional storage utilization.

1. no distributed trace

## experiment linked to the gateway endpoint

- gateway invocation (request, response, token usage)

## experiment for the agent

-- no trace --

2. distributed trace

## experiment linked to the gateway endpoint

- gateway invocation (request, response, token usage)

## experiment for the agent

- agent trace
  - gateway invocation (link to gateway trace, token usage)

How is this PR tested?

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

Does this PR require documentation update?

Does this PR require updating the MLflow Skills repository?

No. You can skip the rest of this section.
Yes. Please link the corresponding PR or explain how you plan to update it.

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)

Enhance the tracing utilities to create distributed spans for gateway calls when a traceparent header is present. This includes new functions to check for traceparent headers, retrieve token usage from spans, and create distributed spans linked to the gateway trace. Update the `maybe_traced_gateway_call` function to accept request headers and integrate tracing into the gateway API invocations. Additionally, add tests to validate the new tracing functionality and ensure proper behavior with and without traceparent headers. Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

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

github-actions · 2026-02-16T10:54:04Z

🛠 DevTools 🛠

Install mlflow from this PR

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

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

Copilot

Pull request overview

Adds distributed tracing support for MLflow Gateway endpoints so incoming requests carrying W3C traceparent context can be linked back to an upstream “agent” trace while preserving a separate gateway trace.

Changes:

Extend maybe_traced_gateway_call to optionally create a distributed span using incoming request headers.
Plumb HTTP request headers through gateway API handlers into the tracing wrapper.
Add a new span attribute key for linking the distributed span to the gateway trace, plus new unit/integration tests.

Reviewed changes

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

File	Description
`mlflow/gateway/tracing_utils.py`	Adds distributed tracing helpers and creates a distributed span when `traceparent` is present.
`mlflow/server/gateway_api.py`	Captures request headers and passes them into `maybe_traced_gateway_call` across gateway endpoints.
`mlflow/tracing/constant.py`	Introduces `SpanAttributeKey.LINKED_GATEWAY_TRACE_ID` for cross-trace linkage.
`tests/gateway/test_tracing_utils.py`	Adds tests for traceparent detection and distributed tracing behavior (streaming + non-streaming).

Comments suppressed due to low confidence (1)

mlflow/gateway/tracing_utils.py:25

_has_traceparent only checks for exact keys "traceparent" and "Traceparent". HTTP header keys are case-insensitive and some frameworks may produce other casings (e.g., "TRACEPARENT"), which would silently disable distributed tracing. Consider normalizing keys (e.g., casefold/lower) before checking.

    if not span:
        return

mlflow/gateway/tracing_utils.py

mlflow/server/gateway_api.py

tests/gateway/test_tracing_utils.py

mlflow/gateway/tracing_utils.py

github-actions · 2026-02-16T11:04:33Z

Documentation preview for f7002d6 is available at:

https://pr-20864--mlflow-docs-preview.netlify.app/docs/latest/

Changed Pages (2)

genai/tracing/app-instrumentation/distributed-tracing (modified)
genai/tracing/integrations/listing/mlflow-ai-gateway (modified)

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.

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

…eval - Introduced a new dataclass `_ProviderSpanInfo` to encapsulate span name and attributes. - Updated `_get_provider_span_attributes` to `_get_provider_span_info`, returning a list of `_ProviderSpanInfo` instead of a dictionary. - Modified the `maybe_create_distributed_span` function to create individual spans for each provider span, improving traceability in distributed tracing scenarios. - Updated tests to reflect changes in the span information retrieval and ensure correct functionality. Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

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

- Extract only W3C trace context headers (traceparent, tracestate) before passing to distributed tracing, avoiding propagation of sensitive headers - Make gateway_experiment_id test fixture resilient to duplicate experiment names by reusing existing experiments Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

TomeHirata · 2026-02-17T09:22:29Z

mlflow/gateway/tracing_utils.py

+
+    try:
+        with set_tracing_context_from_http_request_headers(request_headers):
+            with mlflow.start_span(


I'm actually not sure if there's any clearer way to copy spans with some modification to a distributed span.

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

mlflow/gateway/tracing_utils.py

B-Step62 · 2026-02-19T05:50:02Z

mlflow/gateway/tracing_utils.py

+    return {
+        "endpoint_id": endpoint_config.endpoint_id,
+        "endpoint_name": endpoint_config.endpoint_name,
+    }


Can we log the server URL as well?

The actual server hostname?

mlflow/gateway/tracing_utils.py

B-Step62

Looks good once #20864 (comment) is addressed

B-Step62 · 2026-02-19T05:58:27Z

docs/docs/genai/tracing/integrations/listing/mlflow-ai-gateway.mdx

[Follow-up] I think it is useful to have a guide page about how to use Gateway with different llm/agent libraries, like OpenAI client, LangChain/Graph, LiteLLM, etc, including the distributed tracing setup.

- Add ASCII diagram documenting distributed trace shape - Log server URL (from Host header) in gateway span attributes - Rename provider → model in span naming to align with endpoint IA - Propagate span status to distributed model spans - Add timing assertions in tests 🤖 Generated with Claude Code Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

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

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>

TomeHirata added 2 commits February 16, 2026 19:51

lint

e855b9c

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

Copilot AI review requested due to automatic review settings February 16, 2026 10:53

github-actions bot added size/L Large PR (200-499 LoC) v3.10.0 rn/feature Mention under Features in Changelogs. labels Feb 16, 2026

Copilot started reviewing on behalf of TomeHirata February 16, 2026 10:54 View session

Merge branch 'master' of github.com:mlflow/mlflow into gateway/usage/…

02e9ee1

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

Copilot AI reviewed Feb 16, 2026

View reviewed changes

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

mlflow/server/gateway_api.py Show resolved Hide resolved

tests/gateway/test_tracing_utils.py Outdated Show resolved Hide resolved

mlflow/gateway/tracing_utils.py Show resolved Hide resolved

TomeHirata added 2 commits February 16, 2026 20:04

doc update

8f35350

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

TomeHirata changed the title ~~[draft] Distributed tracing for gateway endpoints~~ Distributed tracing for gateway endpoints Feb 17, 2026

TomeHirata and others added 2 commits February 17, 2026 18:06

comments

6d435b1

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

TomeHirata requested review from B-Step62 and harupy February 17, 2026 09:15

TomeHirata commented Feb 17, 2026

View reviewed changes

TomeHirata added 2 commits February 17, 2026 18:23

copy timstamp

47b72b6

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

lint

e7424d1

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

B-Step62 reviewed Feb 19, 2026

View reviewed changes

B-Step62 approved these changes Feb 19, 2026

View reviewed changes

B-Step62 reviewed Feb 19, 2026

View reviewed changes

TomeHirata and others added 2 commits February 19, 2026 15:06

Merge branch 'master' of github.com:mlflow/mlflow into gateway/usage/…

f7002d6

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

TomeHirata enabled auto-merge February 19, 2026 06:26

github-actions bot assigned B-Step62 Feb 19, 2026

TomeHirata disabled auto-merge February 19, 2026 07:05

TomeHirata merged commit fea5bbc into mlflow:master Feb 19, 2026
56 of 58 checks passed

github-actions bot added size/XL Extra-large PR (500+ LoC) and removed size/L Large PR (200-499 LoC) labels Feb 19, 2026

daniellok-db pushed a commit to daniellok-db/mlflow that referenced this pull request Feb 20, 2026

Distributed tracing for gateway endpoints (mlflow#20864)

c97e1ce

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>

daniellok-db pushed a commit that referenced this pull request Feb 20, 2026

Distributed tracing for gateway endpoints (#20864)

8d4dec9

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Distributed tracing for gateway endpoints#20864

Distributed tracing for gateway endpoints#20864
TomeHirata merged 11 commits intomlflow:masterfrom
TomeHirata:gateway/usage/distributed-tracing

TomeHirata commented Feb 16, 2026 •

edited

Loading

Uh oh!

github-actions bot commented Feb 16, 2026

Install mlflow from this PR

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Feb 16, 2026 •

edited

Loading

Uh oh!

TomeHirata Feb 17, 2026

Uh oh!

Uh oh!

B-Step62 Feb 19, 2026

Uh oh!

TomeHirata Feb 19, 2026

Uh oh!

Uh oh!

Uh oh!

B-Step62 left a comment

Uh oh!

B-Step62 Feb 19, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

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

Does this PR require updating the MLflow Skills repository?

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 16, 2026

Install mlflow from this PR

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 Feb 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

TomeHirata Feb 17, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

B-Step62 Feb 19, 2026

Choose a reason for hiding this comment

Uh oh!

TomeHirata Feb 19, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

B-Step62 left a comment

Choose a reason for hiding this comment

Uh oh!

B-Step62 Feb 19, 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!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

TomeHirata commented Feb 16, 2026 •

edited

Loading

github-actions bot commented Feb 16, 2026 •

edited

Loading

B-Step62 Feb 19, 2026 •

edited

Loading