Repository navigation
Feat/clear mlflow traces pypi 2.0.2 - #67
Conversation
ruivieira
left a comment
ruivieira
left a comment
There was a problem hiding this comment.
LGTM, thanks @suppathak. I've left a couple of questions (inlined).
| "status": {"status_code": info.status, "description": ""}, | ||
| "start_time_unix_nano": s.start_time_ns, | ||
| "end_time_unix_nano": s.end_time_ns, | ||
| "duration_ms": None, |
There was a problem hiding this comment.
@suppathak is the duration available in the traces (e.g. (s.end_time_ns - s.start_time_ns) / 1_000_000) or is it expected to always be None?
| # EVALHUB_MLFLOW_BACKEND=upstream and: pip install mlflow-skinny | ||
| # Required for fetching full traces (with spans) from MLflow when | ||
| # parameters.mlflow_traces_experiment_name is set. | ||
| mlflow-skinny>=2.0.0 |
There was a problem hiding this comment.
@suppathak is this still necessary after the SDK fixes (eval-hub/eval-hub-sdk#134)?
- Replace mlflow-skinny-based _fetch_mlflow_traces_to_dir with sdk client (self.mlflow); drops mlflow-skinny dependency from requirements.txt - Fix duration_ms: calculate from span start/end ns instead of None - Update CLEAR to main_with_sparc branch with [tool-calls] extra - Wire separate_tools parameter into agentic config (SPARC mode) - Update meta/job.json and docs to document separate_tools option
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds MLflow trace ingestion for CLEAR, updates validation and endpoint config handling, revises job parameters and dependency pins, and refreshes docs for trace sources, MLflow fetch steps, and SPARC tool-call analysis. ChangesCLEAR adapter changes
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
|
Hi @ruivieira, addressed both comments in the latest commit (b0bcac7):
Please approve and merge if everything looks good. Thanks! |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
adapters/clear/requirements.txt (1)
3-3: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winGit-branch dependency contradicts the stated "stable PyPI release" goal and is unpinned to a commit.
PR objective
#1says clear-eval was switched from a git ref to "the stable 2.0.2 PyPI release," but this line still pins a mutable git branch (main_with_sparc), not a tag/commit SHA. This also conflicts withmain.py's ImportError text ("pip install clear-eval==2.0.2") and03-local-run.md's claim ("installed from PyPI"). A force-push to that branch would silently change the installed code on next rebuild.♻️ Suggested fix
-clear_eval[tool-calls] @ git+https://github.com/IBM/clear.git@main_with_sparc +clear_eval[tool-calls] @ git+https://github.com/IBM/clear.git@<pinned-commit-sha>Align
main.py's ImportError text and03-local-run.mdwith whichever source is authoritative once SPARC support lands on PyPI.🤖 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 `@adapters/clear/requirements.txt` at line 3, The dependency in requirements.txt is still using a mutable git branch instead of the intended stable PyPI release, so update the clear_eval[tool-calls] entry to the authoritative 2.0.2 source and avoid any branch-based ref. Also make sure the user-facing guidance in main.py (ImportError text) and 03-local-run.md matches that same source of truth so they do not claim PyPI if the package is still coming from git, or vice versa.
🧹 Nitpick comments (3)
adapters/clear/examples/docs/03-local-run.md (1)
20-27: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueClarify trace-source precedence.
The table lists S3, MLflow, and local path as three parallel sources but doesn't state the precedence when more than one is configured (per
main.py, S3-mounted/test_data//datawins over MLflow, which wins overparameters.data_dir). Worth a one-line note to avoid confusion when a job spec configures more than one.🤖 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 `@adapters/clear/examples/docs/03-local-run.md` around lines 20 - 27, Clarify the trace-source selection order in the local run docs: the table in the docs currently presents S3, MLflow, and local path as equal options, but the runtime logic in main.py uses precedence when multiple sources are set. Update the docs near the trace source table to state that S3-mounted /test_data or /data takes priority over MLflow, and MLflow takes priority over parameters.data_dir, so users understand which source is chosen when more than one is configured.adapters/clear/main.py (2)
754-771: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUnreachable validation branch.
epis computed only insideif config.model.url:, and always falls back toconfig.model.url(truthy in this branch) viaor. Sonot epat line 760 can never beTrue— this inner raise is dead code; the only real "missing endpoint" case is already handled by theelifat line 767.♻️ Suggested simplification
if config.model.url: ep = ( config.parameters.get("endpoint_url") or config.parameters.get("inference_url") or config.model.url ) - if agentic_config.get("inference_backend") == "endpoint" and not ep: - raise ValueError( - "model.url or parameters.endpoint_url (or inference_url) is required " - "when inference_backend is 'endpoint'" - ) agentic_config["endpoint_url"] = ep agentic_config["inference_url"] = ep elif agentic_config.get("inference_backend") == "endpoint":🤖 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 `@adapters/clear/main.py` around lines 754 - 771, The validation inside the config model URL branch is unreachable because ep is always truthy there, so remove the dead `if agentic_config.get("inference_backend") == "endpoint" and not ep` check from the `config.model.url` path in `adapters/clear/main.py`. Keep the existing `elif agentic_config.get("inference_backend") == "endpoint"` guard as the only missing-endpoint validation, and leave the endpoint assignment logic in the surrounding configuration block unchanged.
475-506: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winTemp MLflow trace directory is never cleaned up.
tempfile.mkdtemp(prefix="clear_mlflow_traces_")(line 489) creates a temp dir that's never removed after the job (unlike CLEAR's own intermediate output dirs, cleaned in_cleanup_intermediate_files). In a long-lived adapter process handling many jobs, this leaks disk space.♻️ Suggested cleanup
+import shutil as _shutil # already imported as shutil elsewhere mlflow_traces_dir = tempfile.mkdtemp(prefix="clear_mlflow_traces_") _fetch_mlflow_traces_to_dir(self.mlflow, config.parameters, mlflow_traces_dir) data_dir = mlflow_traces_dirTrack
mlflow_traces_dirandshutil.rmtree(mlflow_traces_dir, ignore_errors=True)in afinallyblock after the pipeline run completes.🤖 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 `@adapters/clear/main.py` around lines 475 - 506, The temporary MLflow traces directory created in the CLEAR adapter is never removed, so it can leak disk space across jobs. Update the flow in main.py around the MLflow fetch path to track the directory returned by tempfile.mkdtemp in mlflow_traces_dir and ensure it is deleted with shutil.rmtree(ignore_errors=True) after the pipeline finishes, ideally in a finally block that covers the run. Make sure the cleanup only runs when _fetch_mlflow_traces_to_dir was used and does not affect non-MLflow input paths.
🤖 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 `@adapters/clear/main.py`:
- Line 52: The install guidance in the CLEAR adapter is stale and points users
to the wrong package variant. Update the remediation text in the
`adapters.clear.main` import/error message so it matches the dependency declared
in `requirements.txt`, including the `clear_eval[tool-calls]` extra and the
git-based source/branch rather than `clear-eval==2.0.2`. Keep the message
aligned with the actual install path users should follow when `main` raises the
ImportError.
- Around line 309-370: The trace fetching helper `_fetch_mlflow_traces_to_dir`
is using MLflow APIs and return shapes that don’t match `MlflowClient`. Update
it to call the client’s actual trace methods (`search_traces()` and
`get_trace()`, or the equivalent wrapper methods if this code intentionally
targets a custom client) instead of accessing `mlflow_client.traces`. Also
replace `trace.data.get("spans", [])` with the proper `TraceData` access pattern
such as `.spans` or `.to_dict()` so the span extraction in this function works
without `AttributeError`.
---
Duplicate comments:
In `@adapters/clear/requirements.txt`:
- Line 3: The dependency in requirements.txt is still using a mutable git branch
instead of the intended stable PyPI release, so update the
clear_eval[tool-calls] entry to the authoritative 2.0.2 source and avoid any
branch-based ref. Also make sure the user-facing guidance in main.py
(ImportError text) and 03-local-run.md matches that same source of truth so they
do not claim PyPI if the package is still coming from git, or vice versa.
---
Nitpick comments:
In `@adapters/clear/examples/docs/03-local-run.md`:
- Around line 20-27: Clarify the trace-source selection order in the local run
docs: the table in the docs currently presents S3, MLflow, and local path as
equal options, but the runtime logic in main.py uses precedence when multiple
sources are set. Update the docs near the trace source table to state that
S3-mounted /test_data or /data takes priority over MLflow, and MLflow takes
priority over parameters.data_dir, so users understand which source is chosen
when more than one is configured.
In `@adapters/clear/main.py`:
- Around line 754-771: The validation inside the config model URL branch is
unreachable because ep is always truthy there, so remove the dead `if
agentic_config.get("inference_backend") == "endpoint" and not ep` check from the
`config.model.url` path in `adapters/clear/main.py`. Keep the existing `elif
agentic_config.get("inference_backend") == "endpoint"` guard as the only
missing-endpoint validation, and leave the endpoint assignment logic in the
surrounding configuration block unchanged.
- Around line 475-506: The temporary MLflow traces directory created in the
CLEAR adapter is never removed, so it can leak disk space across jobs. Update
the flow in main.py around the MLflow fetch path to track the directory returned
by tempfile.mkdtemp in mlflow_traces_dir and ensure it is deleted with
shutil.rmtree(ignore_errors=True) after the pipeline finishes, ideally in a
finally block that covers the run. Make sure the cleanup only runs when
_fetch_mlflow_traces_to_dir was used and does not affect non-MLflow input paths.
🪄 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: 20a3d792-82eb-4d51-a670-0f5d33556cc9
📒 Files selected for processing (5)
adapters/clear/examples/docs/03-local-run.mdadapters/clear/examples/docs/05-benchmarks-and-parameters.mdadapters/clear/main.pyadapters/clear/meta/job.jsonadapters/clear/requirements.txt
| def _fetch_mlflow_traces_to_dir(mlflow_client: Any, parameters: dict[str, Any], output_dir: str) -> None: | ||
| """Fetch full traces (with spans) from MLflow using the eval-hub-sdk client. | ||
|
|
||
| Uses the sdk's TracesNamespace (search + get) to retrieve spans via the MLflow | ||
| 3.x API, then converts each trace to the flat format CLEAR's preprocessor expects: | ||
| {trace_id, experiment_id, state, timestamp_ms, execution_time_ms, tags, spans[]} | ||
| """ | ||
| experiment_name = parameters.get("mlflow_traces_experiment_name") | ||
| experiment_id = parameters.get("mlflow_traces_experiment_id") | ||
| max_results = int(parameters.get("mlflow_traces_max_results", 100)) | ||
| run_id = parameters.get("mlflow_traces_run_id") | ||
|
|
||
| if not experiment_name and not experiment_id: | ||
| raise ValueError( | ||
| "Set parameters.mlflow_traces_experiment_name or mlflow_traces_experiment_id" | ||
| ) | ||
|
|
||
| if not experiment_id and experiment_name: | ||
| exp = mlflow_client.get_experiment_by_name(experiment_name) | ||
| if exp is None: | ||
| raise ValueError(f"MLflow experiment not found: {experiment_name!r}") | ||
| experiment_id = exp.experiment_id | ||
|
|
||
| filter_string = None | ||
| if run_id: | ||
| safe_run_id = run_id.strip().replace("'", "''") | ||
| filter_string = f"attribute.run_id = '{safe_run_id}'" | ||
|
|
||
| all_trace_summaries = [] | ||
| page_token: str | None = None | ||
| while len(all_trace_summaries) < max_results: | ||
| page_size = min(100, max_results - len(all_trace_summaries)) | ||
| traces, page_token = mlflow_client.traces.search( | ||
| experiment_ids=[experiment_id], | ||
| max_results=page_size, | ||
| filter_string=filter_string, | ||
| page_token=page_token, | ||
| ) | ||
| if not traces: | ||
| break | ||
| all_trace_summaries.extend(traces) | ||
| if not page_token: | ||
| break | ||
|
|
||
| if not all_trace_summaries: | ||
| raise ValueError( | ||
| f"No traces found in MLflow experiment {experiment_name or experiment_id!r}" | ||
| ) | ||
|
|
||
| Path(output_dir).mkdir(parents=True, exist_ok=True) | ||
| for summary in all_trace_summaries: | ||
| # Re-fetch with get() to include full span data (3.x API) | ||
| trace = mlflow_client.traces.get( | ||
| request_id=summary.info.request_id, | ||
| experiment_id=experiment_id, | ||
| ) | ||
| info = trace.info | ||
| rid = info.request_id | ||
|
|
||
| spans = [] | ||
| for s in trace.data.get("spans", []): | ||
| attrs = _proto_attrs_to_dict(s.get("attributes", [])) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Locate self.mlflow's definition/type and confirm the traces client shape used
rg -n -B3 -A5 'self\.mlflow\s*=' --type=py
rg -n 'class FrameworkAdapter' -A 30 --type=py
rg -n 'mlflow-skinny' -R --type=txt -g '*requirements*.txt'Repository: eval-hub/eval-hub-contrib
Length of output: 2002
🌐 Web query:
MLflow MlflowClient search_traces get_trace TraceData spans to_dict traces namespace documentation
💡 Result:
In MLflow, searching and managing traces is primarily handled through the MlflowClient or the fluent mlflow API [1]. Below is the synthesis of the requested components based on current documentation. Search Traces The MlflowClient.search_traces method allows for programmatic retrieval of traces [2][1]. It supports pagination, which is essential for handling large result sets [2]. Recent updates have introduced a locations parameter (replacing the deprecated experiment_ids) for specifying where to search [3][4]. Fluent API: mlflow.search_traces The fluent API returns a pandas DataFrame by default (if installed), which is suitable for analysis [2][5]. As of MLflow 2.21.1+, you can use the return_type parameter (e.g., return_type="list") to retrieve a list of Trace objects instead [6][2]. Accessing Trace Data Once you have a Trace object (e.g., via get_trace or search_traces), you can interact with its internal structure [7][5]. - TraceData: Each Trace object contains a data attribute, which is an instance of TraceData, holding the collection of spans [8][9]. - Spans: Spans represent individual operations within a trace [9]. You can access them via the trace.data.spans property [8][9]. - Conversion to Dictionary: Trace objects and their components provide a to_dict method to serialize them into dictionary format for easier inspection or export [8][10][9][11]. - trace.to_dict: Serializes the entire trace [9][11]. - span.to_dict: Serializes an individual span [8][9]. Namespace There is no prominent "namespace" parameter documented in the core search_traces or tracing API functions. Filtering is primarily achieved through the filter_string parameter, which utilizes a SQL-like Domain-Specific Language (DSL) to query attributes, tags, and metadata (e.g., filter_string="status = 'OK'") [12][6][2]. Example: Pagination with MlflowClient from mlflow import MlflowClient client = MlflowClient page_token = None all_traces = [] while True: results = client.search_traces( locations=["1"], # Replaces experiment_ids filter_string="status = 'OK'", page_token=page_token,) all_traces.extend(results) if not results.token: break page_token = results.token, top_results=[12][2][8][9][3])
Citations:
- 1: https://mlflow.org/docs/2.20.3/llms/tracing/search-traces.html
- 2: https://mlflow.org/docs/2.22.0/tracing/api/search
- 3: Replace deprecated
experiment_idswithlocationsinsearch_tracescalls mlflow/mlflow#19791 - 4: https://mlflow.org/docs/latest/api_reference/java_api/org/mlflow/api/proto/Service.SearchTracesV3.html
- 5: https://mlflow.org/docs/3.6.0rc0/genai/tracing/observe-with-traces/query-via-sdk/
- 6: https://mlflow.org/docs/3.5.1/genai/tracing/search-traces/
- 7: https://mlflow.org/docs/latest/api_reference/_modules/mlflow/tracing/fluent.html
- 8: https://mlflow.org/docs/latest/api_reference/_modules/mlflow/entities/trace_data.html
- 9: https://docs.databricks.com/gcp/en/mlflow3/genai/tracing/observe-with-traces/access-trace-data
- 10: https://mlflow.org/docs/latest/api_reference/_modules/mlflow/entities/trace.html
- 11: https://mlflow.org/docs/3.0.1/api_reference/_modules/mlflow/entities/trace.html
- 12: https://www.mlflow.org/docs/latest/genai/tracing/search-traces.md
Align the trace fetch logic with MLflow’s actual return types. MlflowClient uses search_traces()/get_trace() (no .traces namespace), and TraceData exposes .spans/.to_dict() rather than .get(). This path will raise AttributeError unless mlflow_client is a custom wrapper that mimics both APIs.
🤖 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 `@adapters/clear/main.py` around lines 309 - 370, The trace fetching helper
`_fetch_mlflow_traces_to_dir` is using MLflow APIs and return shapes that don’t
match `MlflowClient`. Update it to call the client’s actual trace methods
(`search_traces()` and `get_trace()`, or the equivalent wrapper methods if this
code intentionally targets a custom client) instead of accessing
`mlflow_client.traces`. Also replace `trace.data.get("spans", [])` with the
proper `TraceData` access pattern such as `.spans` or `.to_dict()` so the span
extraction in this function works without `AttributeError`.
2 participants
What and why
Three improvements to the IBM CLEAR adapter:
PyPI release — switches
clear-evalfrom the2.0.0-rc.2Git refto the stable
2.0.2PyPI release.MLflow trace fetch — adds the ability to fetch input traces directly
from an MLflow tracking server via
parameters.mlflow_traces_experiment_name.Uses the native
mlflowclient (mlflow-skinny) since the eval-hub-sdk'sTracesNamespacehas a compatibility gap with MLflow 3.x(tracked in eval-hub/eval-hub-sdk#134).
Unified
endpoint_url—model.urlis now passed asendpoint_urlin the CLEAR config for both
litellmandendpointbackends,as suggested by @lilacheden.
Type
Testing
Tested end-to-end with:
inference_backend: litellm+ local traces → exit 0, score 0.73inference_backend: endpoint+ local traces → exit 0, score 0.75Breaking changes
mlflow-skinny>=2.0.0is now a required dependency. Runpip install -r requirements.txtto pick it up.Summary by CodeRabbit