Add safe attribute capture for pydantic_ai by BenWilson2 · Pull Request #19219 · mlflow/mlflow

BenWilson2 · 2025-12-04T18:27:24Z

🛠 DevTools 🛠

Install mlflow from this PR

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

Related Issues/PRs

Resolve #19195

What changes are proposed in this pull request?

Adds allowlist capture of object attributes for auto tracing for pydantic-ai. Our current implementation captures internal state of objects involved in patching async requests which causes pydantic ai's Async wrapper for clients to throw Exceptions (which are ignored, but is still clearly messing with Python GC of these async futures).

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)

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

Copilot

Pull request overview

This PR adds safe attribute capture for pydantic_ai autologging to prevent interference with async cleanup processes. The changes implement allowlist-based attribute extraction instead of capturing all object attributes, which was causing issues with internal async client state management.

Key Changes:

Introduced allowlist-based attribute capture using frozen sets for agent, model, tool, and MCP server attributes
Added helper functions _is_safe_for_serialization and _safe_get_attribute to safely extract and validate attributes
Refactored attribute getter functions to use allowlists instead of iterating over all __dict__ items

Reviewed changes

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

File	Description
mlflow/pydantic_ai/autolog.py	Implements allowlist-based attribute capture with safe serialization checks to avoid capturing client/provider references that interfere with async cleanup
tests/pydantic_ai/test_pydanticai_tracing.py	Adds comprehensive tests for allowlist behavior and validates that client references are not captured in traces

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2025-12-04T18:31:01Z

mlflow/pydantic_ai/autolog.py

+    if isinstance(value, _SAFE_ATTRIBUTE_TYPES):
+        return True
+    if isinstance(value, dict):
+        return all(_is_safe_for_serialization(v) for v in value.values())


list and tuple are included in _SAFE_ATTRIBUTE_TYPES on line 55, which means any list or tuple would be considered safe regardless of their contents. However, the function recursively validates dict values (lines 63-64) but not list/tuple elements. This inconsistency could allow unsafe objects to be captured if they're contained in a list or tuple. Consider adding recursive validation for list and tuple elements similar to the dict validation.

Suggested change

if isinstance(value, _SAFE_ATTRIBUTE_TYPES):

return True

if isinstance(value, dict):

return all(_is_safe_for_serialization(v) for v in value.values())

if isinstance(value, (str, int, float, bool, type(None))):

return True

if isinstance(value, dict):

return all(_is_safe_for_serialization(v) for v in value.values())

if isinstance(value, (list, tuple)):

return all(_is_safe_for_serialization(v) for v in value)

Copilot · 2025-12-04T18:31:01Z

tests/pydantic_ai/test_pydanticai_tracing.py

+@pytest.mark.parametrize(
+    ("getter_func", "mock_attrs", "expected_attrs", "excluded_attrs"),
+    [
+        (
+            _get_agent_attributes,
+            {"name": "test-agent", "system_prompt": "helpful", "retries": 3, "output_type": str},
+            {"name": "test-agent", "system_prompt": "helpful", "retries": 3, "output_type": "str"},
+            ["_client", "provider", "_internal_state"],
+        ),
+        (
+            _get_model_attributes,
+            {"model_name": "gpt-4", "name": "test-model"},
+            {"model_name": "gpt-4", "name": "test-model"},
+            ["client", "_client", "provider", "api_key", "callbacks"],
+        ),
+        (
+            _get_tool_attributes,
+            {"name": "my_tool", "description": "helpful", "max_retries": 2},
+            {"name": "my_tool", "description": "helpful", "max_retries": 2},
+            ["_internal", "func"],
+        ),
+    ],


The new _get_mcp_server_attributes function is not covered by the parametrized test test_attribute_getter_uses_allowlist. Consider adding a test case for this function similar to the existing test cases for _get_agent_attributes, _get_model_attributes, and _get_tool_attributes to ensure consistent behavior and test coverage.

Copilot · 2025-12-04T18:31:02Z

mlflow/pydantic_ai/autolog.py

+_SAFE_ATTRIBUTE_TYPES = (str, int, float, bool, type(None), list, tuple)
+
+
+def _is_safe_for_serialization(value: Any) -> bool:
+    if value is None:
+        return False


type(None) (which is NoneType) is included in _SAFE_ATTRIBUTE_TYPES, but the function explicitly returns False for None values on line 59-60. This creates contradictory logic: line 61 checks isinstance(value, _SAFE_ATTRIBUTE_TYPES) which would include None, but that check is never reached because None is caught earlier. Either remove type(None) from _SAFE_ATTRIBUTE_TYPES or remove the explicit None check on line 59-60.

github-actions · 2025-12-04T18:35:29Z

Documentation preview for d607320 is available at:

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

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

serena-ruan · 2025-12-08T06:17:30Z

mlflow/pydantic_ai/autolog.py

+# Allowlists for safe attributes to extract from pydantic_ai objects.
+# Using allowlists instead of denylists to avoid capturing client/provider
+# references that can interfere with async cleanup (e.g., httpx client lifecycle).


Could you explain which part causes the error of

Exception ignored in: <function AsyncHttpxClientWrapper.del at ...> AttributeError: 'AsyncHttpxClientWrapper' object has no attribute '_state'

? If we add allowlist does that mean in the future we need to extend the list to support new attributes (which seems more maintenance burden)?

The fact that we're pulling in all attributes is the part that's causing an issue. We're pulling in the base class attributes from within their library which is causing a reference to exist when their async loop manager attempts to GC the object which raises the exception of a reference not being present.
While I agree this is a bit of an addition of maintenance burden, the fact that we're doing the 'easier to maintain' approach and that is causing exceptions to be raised in client code is not good.

What about skipping all attributes starting with underscores?

Let me give that a shot :) I think that might be a great compromise between maintainability and correctness :D Thanks @serena-ruan !

serena-ruan · 2025-12-09T06:28:23Z

mlflow/pydantic_ai/autolog.py

+    if hasattr(value, "__dataclass_fields__"):
+        return True


Is this always true?

There's a cleaner way to do this... instead of using dunder commands I'll use the standard method just as a guard to protect unserializable subclass components from the library

serena-ruan · 2025-12-09T06:58:55Z

mlflow/pydantic_ai/autolog.py

+# Allowlists for safe attributes to extract from pydantic_ai objects.
+# Using allowlists instead of denylists to avoid capturing client/provider
+# references that can interfere with async cleanup (e.g., httpx client lifecycle).


What about skipping all attributes starting with underscores?

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

serena-ruan

LGTM!

Add safe attribute capture for pydantic_ai

a37a972

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

Copilot AI review requested due to automatic review settings December 4, 2025 18:27

github-actions bot added area/tracing MLflow Tracing and its integrations rn/bug-fix Mention under Bug Fixes in Changelogs. v3.7.0 labels Dec 4, 2025

Copilot started reviewing on behalf of BenWilson2 December 4, 2025 18:27 View session

Copilot finished reviewing on behalf of BenWilson2 December 4, 2025 18:30

Copilot AI reviewed Dec 4, 2025

View reviewed changes

BenWilson2 added v3.7.1 and removed v3.7.0 labels Dec 4, 2025

BenWilson2 and others added 2 commits December 4, 2025 22:15

feedback

98a4b63

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

Merge branch 'master' into pydantic-ai-safe-capture

184c863

BenWilson2 added the team-review Trigger a team review request label Dec 5, 2025

github-actions bot requested review from B-Step62, TomeHirata, WeichenXu123, daniellok-db, harupy, kevin-lyn, serena-ruan and xq-yin December 5, 2025 04:59

serena-ruan reviewed Dec 8, 2025

View reviewed changes

serena-ruan reviewed Dec 9, 2025

View reviewed changes

github-actions bot assigned serena-ruan Dec 9, 2025

BenWilson2 requested a review from serena-ruan December 10, 2025 00:57

BenWilson2 added 2 commits December 9, 2025 19:57

apply suggestions

16dfcf6

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

fix test bug

d607320

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

serena-ruan approved these changes Dec 11, 2025

View reviewed changes

BenWilson2 added this pull request to the merge queue Dec 11, 2025

github-merge-queue bot removed this pull request from the merge queue due to failed status checks Dec 11, 2025

BenWilson2 added this pull request to the merge queue Dec 11, 2025

Merged via the queue into mlflow:master with commit 4dfd71d Dec 11, 2025
84 of 86 checks passed

BenWilson2 deleted the pydantic-ai-safe-capture branch December 11, 2025 16:20

MarkVasile mentioned this pull request Mar 7, 2026

Exclude run_config from span attributes #21454

Open

16 tasks

Conversation

BenWilson2 commented Dec 4, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Install mlflow from this 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!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Dec 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

serena-ruan Dec 8, 2025

Choose a reason for hiding this comment

Uh oh!

BenWilson2 Dec 8, 2025

Choose a reason for hiding this comment

Uh oh!

serena-ruan Dec 9, 2025

Choose a reason for hiding this comment

Uh oh!

BenWilson2 Dec 10, 2025

Choose a reason for hiding this comment

Uh oh!

serena-ruan Dec 9, 2025

Choose a reason for hiding this comment

Uh oh!

BenWilson2 Dec 10, 2025

Choose a reason for hiding this comment

Uh oh!

serena-ruan Dec 9, 2025

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!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

BenWilson2 commented Dec 4, 2025 •

edited by github-actions bot

Loading

github-actions bot commented Dec 4, 2025 •

edited

Loading