get_trace support for stores by serena-ruan · Pull Request #18556 · mlflow/mlflow

serena-ruan · 2025-10-28T10:59:19Z

🥞 Stacked PR

Use this link to review incremental changes.

stack/get_trace_stores [Files changed]
- stack/get_trace_client [Files changed]

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

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 · 2025-11-10T06:56:18Z

Documentation preview for 5d8cc32 is available at:

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

serena-ruan · 2025-11-11T11:06:25Z

mlflow/protos/service.proto

-  // Whether to allow partial traces. Default to True.
-  optional bool allow_partial = 2 [default = true];;
+  // Whether to allow partial traces. Default to False.
+  optional bool allow_partial = 2 [default = false];;


Updating the default value to False since after second thought I think that the only case that users may want to fetch in-progress trace is through MLflow UI. There's no specific reason user may want to get in-progress traces using SDK, and mlflow.get_trace is depended by lots of places like eval so we should retry to fetch the trace by default.

harupy · 2025-11-12T01:02:14Z

~~Running into a network issue. Cannot join the standup.~~

serena-ruan · 2025-11-12T01:10:06Z

Running into a network issue. Cannot join the standup

@harupy Haha we're not having standup today, you can leave async updates in slack afterwards

mlflow/protos/service.proto

mlflow/server/handlers.py

mlflow/store/tracking/abstract_store.py

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

mlflow/server/handlers.py

mlflow/store/tracking/abstract_store.py

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

mlflow/store/tracking/sqlalchemy_store.py

harupy · 2025-11-13T10:18:14Z

/review

🚀 Review workflow started

✅ Review completed.

Review Output

PR Review Complete

I've reviewed PR #18556 "get_trace support for stores" and found 6 issues that need attention:

Critical Issues (Bugs):

Missing time import in sqlalchemy_store.py (line 3548) - This will cause a runtime NameError
Incorrect return type annotation (sqlalchemy_store.py:3542) - The function is annotated as Trace | None but should be Trace since it raises an exception instead of returning None for missing traces
Potential bug with None return (sqlalchemy_store.py:3542-3550) - When allow_partial=False and the trace isn't fully exported after retries, the function returns None, which will cause an exception when the handler tries to call .to_proto() on it

Style Guide Violations:

Try-catch scope too broad (handlers.py:3167-3169) - The safe operation trace_data = trace.data.to_dict() should be moved outside the try block or into an else block
Try-catch scope too broad (handlers.py:3176-3183) - Similar issue with exception handling scope
Mock setup pattern (test_handlers.py:2006+) - While this is a minor style preference and may not be applicable due to fixture usage, the guide recommends setting return_value in the patch declaration when possible

The most critical issues are #1 (missing import) and #3 (None return bug) as they will cause runtime failures. I've added detailed review comments with suggestions for each issue on the PR.

mlflow/server/handlers.py

mlflow/store/tracking/sqlalchemy_store.py

tests/server/test_handlers.py

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

tests/server/test_handlers.py

harupy

Left a few more comments, otherwise LGTM :)

mlflow/server/handlers.py

mlflow/store/tracking/abstract_store.py

mlflow/server/handlers.py

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

harupy · 2025-11-18T11:36:01Z

mlflow/server/handlers.py

+        # allow partial so the front end can render in-progress traces
+        trace = _get_tracking_store().get_trace(request_id, allow_partial=True)
+    except MlflowTracingException:
+        pass
+    except MlflowNotImplementedException:
+        # default to batch get traces if get_trace is not implemented
+        try:
+            traces = _get_tracking_store().batch_get_traces([request_id], None)
+        # For stores that don't support batch get traces, or if the trace data is not stored in the
+        # tracking store, we need to get the trace data from the artifact repository.
+        except (MlflowTracingException, MlflowNotImplementedException):
+            pass
+        else:
+            if len(traces) != 1:
+                raise MlflowException(
+                    f"Trace with id={request_id} not found.",
+                    error_code=RESOURCE_DOES_NOT_EXIST,
+                )
+            trace_data = traces[0].data.to_dict()
+    else:
+        trace_data = trace.data.to_dict()


A helper function doesn't help simplify the code?

Emm we can move this part to a helper function, but I don't think it'll be simpler.

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

serena-ruan mentioned this pull request Oct 28, 2025

Add GetTrace proto #18555

Merged

29 tasks

serena-ruan force-pushed the stack/get_trace_stores branch 2 times, most recently from 6743be6 to 51e815f Compare November 10, 2025 06:40

serena-ruan marked this pull request as ready for review November 10, 2025 06:43

github-actions bot added the rn/none List under Small Changes in Changelogs. label Nov 10, 2025

serena-ruan force-pushed the stack/get_trace_stores branch 2 times, most recently from 41211da to e023a53 Compare November 11, 2025 10:22

serena-ruan mentioned this pull request Nov 11, 2025

Update client for get_trace #18778

Merged

29 tasks

serena-ruan commented Nov 11, 2025

View reviewed changes

serena-ruan added the team-review Trigger a team review request label Nov 11, 2025

github-actions bot requested review from B-Step62, BenWilson2, TomeHirata, WeichenXu123, daniellok-db, harupy, kevin-lyn and xq-yin November 11, 2025 12:09

harupy reviewed Nov 12, 2025

View reviewed changes

mlflow/protos/service.proto Outdated Show resolved Hide resolved

harupy reviewed Nov 12, 2025

View reviewed changes

mlflow/server/handlers.py Show resolved Hide resolved

harupy reviewed Nov 12, 2025

View reviewed changes

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

serena-ruan added 3 commits November 13, 2025 11:35

get_trace suport for stores

5a4f650

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

update

1906cfe

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

fix and comments

f074bcd

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

serena-ruan force-pushed the stack/get_trace_stores branch from 7edfb9e to f074bcd Compare November 13, 2025 03:38

lint

c3a251a

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

harupy reviewed Nov 13, 2025

View reviewed changes

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

harupy reviewed Nov 13, 2025

View reviewed changes

mlflow/store/tracking/abstract_store.py Show resolved Hide resolved

harupy reviewed Nov 13, 2025

View reviewed changes

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

comments

8138c7b

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

harupy reviewed Nov 13, 2025

View reviewed changes

mlflow/store/tracking/sqlalchemy_store.py Outdated Show resolved Hide resolved

github-actions bot reviewed Nov 13, 2025

View reviewed changes

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

github-actions bot reviewed Nov 13, 2025

View reviewed changes

mlflow/server/handlers.py Show resolved Hide resolved

github-actions bot reviewed Nov 13, 2025

View reviewed changes

mlflow/store/tracking/sqlalchemy_store.py Outdated Show resolved Hide resolved

github-actions bot reviewed Nov 13, 2025

View reviewed changes

mlflow/store/tracking/sqlalchemy_store.py Outdated Show resolved Hide resolved

github-actions bot reviewed Nov 13, 2025

View reviewed changes

mlflow/store/tracking/sqlalchemy_store.py Outdated Show resolved Hide resolved

github-actions bot reviewed Nov 13, 2025

View reviewed changes

tests/server/test_handlers.py Show resolved Hide resolved

serena-ruan added 2 commits November 14, 2025 15:38

comments

627ad07

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

type hint

2bfa92f

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

serena-ruan requested a review from harupy November 18, 2025 07:05