Fix shared cluster Py4j statefulness issue by BenWilson2 · Pull Request #19139 · mlflow/mlflow

BenWilson2 · 2025-12-01T18:13:44Z

🛠 DevTools 🛠

Install mlflow from this PR

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

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Fix for shared clusters issue with py4j rejected connection timeouts when saving / logging models.

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>

github-actions · 2025-12-01T18:20:47Z

Documentation preview for 3ca9258 is available at:

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

harupy · 2025-12-02T02:38:13Z

tests/utils/test_requirements_utils.py

+        try:
+            _capture_imported_modules("fake/model/path", "pyfunc")
+        except MlflowException:
+            pass


It's unclear whether this should succeed or fail. Let's use pytest.raises.

harupy · 2025-12-02T02:38:25Z

/review

🚀 Review workflow started

✅ Review completed.

Review Output

Perfect! I've completed the PR review. I found one issue that violates the Python style guide:

Review Summary

PR: #19139 - Fix shared cluster Py4j statefulness issue

Issues Found: 1

Issue Details

Missing mock assertion in tests/utils/test_requirements_utils.py:748
- The test mocks _run_command but doesn't verify it was called
- This violates the style guide rule: "Always Verify Mock Calls with Assertions"
- I've added a comment with a suggestion to add mock_run.assert_called_once()

The implementation itself looks good - it correctly removes the PySpark gateway environment variables to prevent subprocess corruption of the parent process's py4j connection. The comments clearly explain why this is necessary.

github-actions · 2025-12-02T02:39:45Z

tests/utils/test_requirements_utils.py

+            "mlflow.utils.requirements_utils._download_artifact_from_uri",
+            return_value=str(tmp_path),
+        ),
+    ):


The mocked _run_command function should have an assertion to verify it was called. Without this, the test could pass even if _run_command is never invoked.

Consider adding an assertion after the try-except block:

Suggested change

):

with (

mock.patch(

"mlflow.utils.requirements_utils._run_command",

side_effect=mock_run_command,

) as mock_run,

mock.patch(

"mlflow.utils.requirements_utils._download_artifact_from_uri",

return_value=str(tmp_path),

),

):

try:

_capture_imported_modules("fake/model/path", "pyfunc")

except MlflowException:

pass

mock_run.assert_called_once()

assert "PYSPARK_GATEWAY_PORT" not in captured_env

assert "PYSPARK_GATEWAY_SECRET" not in captured_env

This follows the Python style guide rule: "Always Verify Mock Calls with Assertions" (see dev/guides/python.md).

🤖 Generated with Claude Code

let's validate that the mock for _download_artifact_from_uri is used too.

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

attempt fix 1

2d35ab4

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

github-actions bot added v3.7.0 area/models MLmodel format, model serialization/deserialization, flavors rn/bug-fix Mention under Bug Fixes in Changelogs. labels Dec 1, 2025

BenWilson2 added 2 commits December 1, 2025 14:45

add test

d56156b

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

add test

28ab7d5

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

BenWilson2 changed the title ~~[ Testing ] Attempt shared cluster py4j state corruption fix~~ Fix shared cluster Py4j statefulness issue Dec 1, 2025

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

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

harupy reviewed Dec 2, 2025

View reviewed changes

github-actions bot reviewed Dec 2, 2025

View reviewed changes

cleanup test

3ca9258

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

BenWilson2 requested a review from harupy December 2, 2025 19:17

harupy approved these changes Dec 3, 2025

View reviewed changes

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

Merged via the queue into mlflow:master with commit f0cd3e9 Dec 3, 2025
66 of 69 checks passed

BenWilson2 deleted the serverless branch December 3, 2025 01:43

BenWilson2 added a commit to BenWilson2/mlflow that referenced this pull request Dec 4, 2025

Fix shared cluster Py4j statefulness issue (mlflow#19139)

bb66b04

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

BenWilson2 added a commit that referenced this pull request Dec 4, 2025

Fix shared cluster Py4j statefulness issue (#19139)

b5bd021

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

TomeHirata mentioned this pull request Feb 20, 2026

Remove PYSPARK_GATEWAY_PORT for transformer trainer to fix py4j issue #21009

Open

31 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix shared cluster Py4j statefulness issue#19139

Fix shared cluster Py4j statefulness issue#19139
BenWilson2 merged 4 commits intomlflow:masterfrom
BenWilson2:serverless

BenWilson2 commented Dec 1, 2025 •

edited by github-actions bot

Loading

Uh oh!

github-actions bot commented Dec 1, 2025 •

edited

Loading

Uh oh!

harupy Dec 2, 2025 •

edited

Loading

Uh oh!

harupy commented Dec 2, 2025 •

edited by github-actions bot

Loading

Review Summary

Issue Details

Uh oh!

github-actions bot Dec 2, 2025

Uh oh!

harupy Dec 2, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

-    ):
+    with (
+        mock.patch(
+            "mlflow.utils.requirements_utils._run_command",
+            side_effect=mock_run_command,
+        ) as mock_run,
+        mock.patch(
+            "mlflow.utils.requirements_utils._download_artifact_from_uri",
+            return_value=str(tmp_path),
+        ),
+    ):
+        try:
+            _capture_imported_modules("fake/model/path", "pyfunc")
+        except MlflowException:
+            pass
+    mock_run.assert_called_once()
+    assert "PYSPARK_GATEWAY_PORT" not in captured_env
+    assert "PYSPARK_GATEWAY_SECRET" not in captured_env

Conversation

BenWilson2 commented Dec 1, 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!

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

Uh oh!

Uh oh!

harupy Dec 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Review Summary

Issue Details

Uh oh!

github-actions bot Dec 2, 2025

Choose a reason for hiding this comment

Uh oh!

harupy Dec 2, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

BenWilson2 commented Dec 1, 2025 •

edited by github-actions bot

Loading

github-actions bot commented Dec 1, 2025 •

edited

Loading

harupy Dec 2, 2025 •

edited

Loading

harupy commented Dec 2, 2025 •

edited by github-actions bot

Loading