feat: Update mlflow adapter to be compatible with latest 3.x versions#144
Merged
Conversation
Contributor
📝 WalkthroughWalkthroughMLflow adapter changes update HTTP error handling, artifact upload path resolution, trace response parsing, and trace retrieval/materialization calls for MLflow 3.x compatibility. ChangesMLflow 3.x Trace API Compatibility
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/evalhub/adapter/mlflow.py (1)
849-878:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHydrate each trace before writing files.
materialize()still serializes the objects returned bysearch(). On the MLflow 3.x path, those rows are metadata-only, sotrace.datastays empty and the exported JSON still lacks spans — the exact behavior issue#134calls out. Fetch eachrequest_idthroughself.get(..., experiment_id)before writing the file.💡 Suggested fix
for trace in traces: + if trace.info.request_id: + trace = self.get(trace.info.request_id, experiment_id) tid = re.sub(r"[^a-zA-Z0-9_\-]", "_", trace.info.request_id) if not tid: tid = uuid.uuid4().hex🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/evalhub/adapter/mlflow.py` around lines 849 - 878, The traces returned by self.search(...) are metadata-only on MLflow 3.x, so before writing each file call self.get(request_id, experiment_id=experiment_id) to hydrate the full trace (replace or merge trace.data with the hydrated result) — update the loop that iterates over traces from search() to fetch the full trace via self.get(trace.info.request_id, experiment_id=experiment_id) and then use that hydrated trace when building trace_dict and writing to file; ensure you still fall back to the uuid when request_id is empty and preserve the existing filename logic (references: search, get, trace.info, trace.data).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unit/test_mlflow_traces.py`:
- Around line 77-132: Add pytest unit markers to the new tests in
tests/unit/test_mlflow_traces.py so they participate in marker-based selection:
either decorate each test function (test_parse_trace_flat_search_response,
test_parse_trace_info_endpoint_response,
test_artifact_server_path_odh_workspace,
test_artifact_server_path_upstream_run_ids,
test_artifact_server_path_http_proxied) with `@pytest.mark.unit`, or set
pytestmark = [pytest.mark.unit] at the module top; ensure pytest is imported if
adding decorators.
---
Outside diff comments:
In `@src/evalhub/adapter/mlflow.py`:
- Around line 849-878: The traces returned by self.search(...) are metadata-only
on MLflow 3.x, so before writing each file call self.get(request_id,
experiment_id=experiment_id) to hydrate the full trace (replace or merge
trace.data with the hydrated result) — update the loop that iterates over traces
from search() to fetch the full trace via self.get(trace.info.request_id,
experiment_id=experiment_id) and then use that hydrated trace when building
trace_dict and writing to file; ensure you still fall back to the uuid when
request_id is empty and preserve the existing filename logic (references:
search, get, trace.info, trace.data).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 88318ad6-77ae-437c-8207-4e309ad86bf9
📒 Files selected for processing (2)
src/evalhub/adapter/mlflow.pytests/unit/test_mlflow_traces.py
Comment on lines
+77
to
+132
| def test_parse_trace_flat_search_response() -> None: | ||
| """Trace search returns trace fields at the top level, not wrapped in ``info``.""" | ||
| raw = { | ||
| "request_id": "tr-abc123", | ||
| "experiment_id": "1", | ||
| "timestamp_ms": 1700000000000, | ||
| "execution_time_ms": 45, | ||
| "status": "OK", | ||
| "tags": [{"key": "framework", "value": "langgraph"}], | ||
| "request_metadata": [{"key": "source", "value": "test"}], | ||
| } | ||
| trace = _parse_trace(raw) | ||
| assert trace.info.request_id == "tr-abc123" | ||
| assert trace.info.experiment_id == "1" | ||
| assert trace.info.status == "OK" | ||
| assert trace.info.tags == {"framework": "langgraph"} | ||
| assert trace.info.request_metadata == {"source": "test"} | ||
|
|
||
|
|
||
| def test_parse_trace_info_endpoint_response() -> None: | ||
| raw = { | ||
| "trace_info": { | ||
| "request_id": "tr-xyz", | ||
| "experiment_id": "2", | ||
| "timestamp_ms": 100, | ||
| "execution_time_ms": 10, | ||
| "status": "OK", | ||
| "tags": [], | ||
| "request_metadata": [], | ||
| } | ||
| } | ||
| trace = _parse_trace(raw) | ||
| assert trace.info.request_id == "tr-xyz" | ||
| assert trace.info.experiment_id == "2" | ||
|
|
||
|
|
||
| def test_artifact_server_path_odh_workspace() -> None: | ||
| uri = "mlflow-artifacts:/workspaces/ws/1/run-abc/artifacts" | ||
| path = MlflowClient._artifact_server_path(uri, "results/out.json") | ||
| assert path == "/api/2.0/mlflow-artifacts/artifacts/1/run-abc/artifacts/results/out.json" | ||
|
|
||
|
|
||
| def test_artifact_server_path_upstream_run_ids() -> None: | ||
| uri = "/private/tmp/mlflow/artifacts/6/run-abc/artifacts" | ||
| path = MlflowClient._artifact_server_path( | ||
| uri, "results/out.json", experiment_id="6", run_id="run-abc" | ||
| ) | ||
| assert path == ( | ||
| "/api/2.0/mlflow-artifacts/artifacts/6/run-abc/artifacts/results/out.json" | ||
| ) | ||
|
|
||
|
|
||
| def test_artifact_server_path_http_proxied() -> None: | ||
| uri = "http://localhost:5000/api/2.0/mlflow-artifacts/artifacts/1/run-abc/artifacts" | ||
| path = MlflowClient._artifact_server_path(uri, "out.json") | ||
| assert path == "/api/2.0/mlflow-artifacts/artifacts/1/run-abc/artifacts/out.json" |
Contributor
There was a problem hiding this comment.
Add the required pytest markers to these new tests.
These additions are unmarked, so they will not participate in the repo's marker-based test selection. Add @pytest.mark.unit here, or set pytestmark once at module scope if this whole file is a unit suite.
💡 Suggested fix
import pytest
from evalhub.adapter.mlflow import (
MlflowClient,
TracesNamespace,
_parse_trace,
)
+
+pytestmark = [pytest.mark.unit]As per coding guidelines, tests/**/*.py: Mark tests with pytest markers: @pytest.mark.unit, @pytest.mark.integration, @pytest.mark.adapter, @pytest.mark.e2e.
Also applies to: 138-160
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/test_mlflow_traces.py` around lines 77 - 132, Add pytest unit
markers to the new tests in tests/unit/test_mlflow_traces.py so they participate
in marker-based selection: either decorate each test function
(test_parse_trace_flat_search_response, test_parse_trace_info_endpoint_response,
test_artifact_server_path_odh_workspace,
test_artifact_server_path_upstream_run_ids,
test_artifact_server_path_http_proxied) with `@pytest.mark.unit`, or set
pytestmark = [pytest.mark.unit] at the module top; ensure pytest is imported if
adding decorators.
Source: Coding guidelines
ruivieira
changed the title
ruivieira
approved these changes
Jun 29, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What and why
Updates the MLflow adapter to be compatible with 3.x MLflow servers
Closes #134
Type
Testing
cc @ruivieira @suppathak
Summary by CodeRabbit
Summary by CodeRabbit
Release Notes
New Features
Bug Fixes
Tests