Bug fix: TraceStatus hydration from proto #12044

mparkhe · 2024-05-19T01:58:15Z

🛠 DevTools 🛠

Install mlflow from this PR

pip install git+https://github.com/mlflow/mlflow.git@refs/pull/12044/merge

Checkout with GitHub CLI

gh pr checkout 12044

What changes are proposed in this pull request?

Fixing a bug in trace info deserialization

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

Interface

area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
area/windows: Windows support

Language

language/r: R APIs and clients
language/java: Java APIs and clients
language/new: Proposals for new client languages

Integrations

integrations/azure: Azure and Azure ML integrations
integrations/sagemaker: SageMaker integrations
integrations/databricks: Databricks integrations

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 · 2024-05-19T01:58:36Z

Documentation preview for 968eb21 will be available when this CircleCI job
completes successfully.

More info

Ignore this comment if this PR does not change the documentation.
It takes a few minutes for the preview to be available.
The preview is updated when a new commit is pushed to this PR.
This comment was created by https://github.com/mlflow/mlflow/actions/runs/9145908554.

github-actions · 2024-05-19T01:58:40Z

@mparkhe Thank you for the contribution! Could you fix the following issue(s)?

⚠ DCO check

The DCO check failed. Please sign off your commit(s) by following the instructions here. See https://github.com/mlflow/mlflow/blob/master/CONTRIBUTING.md#sign-your-work for more details.

Signed-off-by: Mani Parkhe <mani@databricks.com>

dbczumar · 2024-05-19T04:47:58Z

mlflow/entities/trace_info.py

@@ -83,7 +83,7 @@ def to_dict(self):
        Update status field to the string value for serialization.
        """
        trace_info_dict = asdict(self)
-        trace_info_dict["status"] = self.status.value


Following up from our thread offline, the defined type of self.status in TraceInfo is TraceStatus, so the preexisting logic in to_dict() should in theory be correct. My guess is that TraceInfo is being constructed with a string type for TraceStatus (which is the incorrect type) when a trace is read from the backend.

The problem seems to lie in

mlflow/mlflow/store/tracking/rest_store.py

Line 338 in f54297d

return TraceInfo.from_proto(response_proto.trace_info)

->

mlflow/mlflow/entities/trace_info.py

Line 75 in f54297d

status=TraceStatus.from_proto(proto.status),

->

mlflow/mlflow/entities/trace_status.py

Line 24 in f54297d

return ProtoTraceStatus.Name(proto_status)

here, we see that TraceStatus.from_proto() doesn't return a TraceStatus like it should. Instead, it returns a string. We need to fix the implementation of TraceStatus.from_proto() and then add test coverage verifying that a TraceStatus constructed from a proto can be converted to dict / JSON and back.

Yeah! I realized that the fix was superficial and hence incorrect. Fixed the from_proto hydration and added a bunch of tests

dbczumar

Following from https://github.com/mlflow/mlflow/pull/12044/files#r1605945236, I think the fix lies in TraceStatus.from_proto()

Signed-off-by: Mani Parkhe <mani@databricks.com>

dbczumar · 2024-05-19T06:47:39Z

tests/entities/test_trace_info.py

@@ -1,11 +1,14 @@
+import pytest
+


Can we add a test case where we do TraceInfo.from_dict(TraceInfo.from_proto(...).to_dict()) to make sure that loading a trace from proto and converting to dict & back works properly? This will prevent recurrence of the issue.

dbczumar · 2024-05-19T06:48:17Z

tests/entities/test_trace_info.py

+def test_to_dict(trace_info):
+    trace_as_dict = trace_info.to_dict()
+    assert trace_as_dict == {
+        "request_id": "request_id",
+        "experiment_id": "test_experiment",
+        "timestamp_ms": 0,
+        "execution_time_ms": 1,
+        "status": "OK",
+        "request_metadata": {
+            "foo": "bar",
+            "k" * 1000: "v" * 1000,
+        },
+        "tags": {
+            "baz": "qux",
+            "k" * 2000: "v" * 2000,
+        },
+    }


Awesome that we're adding this. I think this still passes on master because the status field of TraceInfo is being set to a TraceStatus object in https://github.com/mlflow/mlflow/pull/12044/files#diff-5aee524692b8b8837b083f08c50a7feb4ebc2d192b97548527ad6978173bc99dR16.

If we also add https://github.com/mlflow/mlflow/pull/12044/files#r1605960450, then we will prevent recurrence.

dbczumar

LGTM once https://github.com/mlflow/mlflow/pull/12044/files#r1605960450 is addressed. Thanks @mparkhe ! :D

Signed-off-by: Mani Parkhe <mani@databricks.com>

Bug fix: Deserialize trace info into python dictionary

2012416

github-actions bot added the patch-2.12.3 label May 19, 2024

github-actions bot added area/tracking Tracking service, tracking client APIs, autologging rn/bug-fix Mention under Bug Fixes in Changelogs. labels May 19, 2024

signing PR

d2efa05

Signed-off-by: Mani Parkhe <mani@databricks.com>

dbczumar reviewed May 19, 2024

View reviewed changes

dbczumar requested changes May 19, 2024

View reviewed changes

Fix TraceStatus hydration from proto

dab3219

Signed-off-by: Mani Parkhe <mani@databricks.com>

mparkhe requested a review from dbczumar May 19, 2024 06:39

formatting

da2a57e

Signed-off-by: Mani Parkhe <mani@databricks.com>

mparkhe changed the title ~~Bug fix: Deserialize trace info into python dictionary~~ Bug fix: TraceStatus hydration from proto May 19, 2024

dbczumar reviewed May 19, 2024

View reviewed changes

dbczumar approved these changes May 19, 2024

View reviewed changes

mparkhe added 2 commits May 19, 2024 00:12

added end-to-end de/serialization test for TraceInfo

000413b

Signed-off-by: Mani Parkhe <mani@databricks.com>

Adding inline comments

968eb21

Signed-off-by: Mani Parkhe <mani@databricks.com>

mparkhe requested a review from dbczumar May 19, 2024 07:53

dbczumar merged commit 2c7906a into mlflow:master May 19, 2024
39 of 41 checks passed

mparkhe deleted the trace_info branch May 19, 2024 08:05

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Bug fix: TraceStatus hydration from proto #12044

Bug fix: TraceStatus hydration from proto #12044

mparkhe commented May 19, 2024 •

edited by github-actions bot

github-actions bot commented May 19, 2024 •

edited

github-actions bot commented May 19, 2024

dbczumar May 19, 2024 •

edited

mparkhe May 19, 2024

dbczumar left a comment

dbczumar May 19, 2024

dbczumar May 19, 2024

mparkhe May 19, 2024

dbczumar left a comment

Bug fix: TraceStatus hydration from proto #12044

Bug fix: TraceStatus hydration from proto #12044

Conversation

mparkhe commented May 19, 2024 • edited by github-actions bot

Install mlflow from this PR

Checkout with GitHub CLI

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?

github-actions bot commented May 19, 2024 • edited

github-actions bot commented May 19, 2024

⚠ DCO check

dbczumar May 19, 2024 • edited

Choose a reason for hiding this comment

mparkhe May 19, 2024

Choose a reason for hiding this comment

dbczumar left a comment

Choose a reason for hiding this comment

dbczumar May 19, 2024

Choose a reason for hiding this comment

dbczumar May 19, 2024

Choose a reason for hiding this comment

mparkhe May 19, 2024

Choose a reason for hiding this comment

dbczumar left a comment

Choose a reason for hiding this comment

mparkhe commented May 19, 2024 •

edited by github-actions bot

github-actions bot commented May 19, 2024 •

edited

dbczumar May 19, 2024 •

edited