Work around Dataproc Py4J weak container references [reduced-it] [fast-ut] [databricks] - #15880
Work around Dataproc Py4J weak container references [reduced-it] [fast-ut] [databricks]#15880wjxiz1992 wants to merge 2 commits into
Conversation
Signed-off-by: Allen Xu <allxu@nvidia.com>
wjxiz1992
changed the title
firestarman
left a comment
firestarman
left a comment
There was a problem hiding this comment.
Inline notes on the probe, session-start hook, and regression coverage. The overall approach still looks right for a temporary IT-only workaround.
| container = _ProbeContainer() | ||
| member = java_member_class( | ||
| '_spark_rapids_container_lifetime_probe', container, 'o0', _ProbeGatewayClient()) | ||
| return isinstance(member.container, weakref.ReferenceType) |
There was a problem hiding this comment.
This probe constructs JavaMember with a fake gateway client and is invoked from pytest_sessionstart with no try/except. If Dataproc's patched __init__ touches more of the gateway than upstream, the xdist worker fails before any test runs. Probe failures should log and skip the patch rather than abort session startup.
| import pyspark | ||
| from py4j.java_gateway import java_import | ||
|
|
||
| # Dataproc 2.2.86 can delete temporary JVM targets before chained Py4J calls complete. |
There was a problem hiding this comment.
This is a process-wide monkeypatch of the IT harness, not a Dataproc/Py4J platform fix. That boundary is fine, but a green Linux/Databricks IT run should not be read as fixing user pyspark jobs on 2.2.86. Also worth wrapping this call so a probe exception cannot take down the whole worker.
| logging.warning( | ||
| "Detected weak Py4J JavaMember container references; applying the temporary " | ||
| "strong-reference workaround for https://github.com/NVIDIA/cudf-spark/issues/15805") | ||
| return True |
| def test_py4j_chained_temporary_container_call(): | ||
| spark = get_spark_i_know_what_i_am_doing() | ||
|
|
||
| assert spark._jvm.java.util.Collections.emptyList().size() == 0 |
There was a problem hiding this comment.
The Dataproc failures were chained lookups such as scala_map.get(key).get() in spark_session.py. emptyList().size() is always green on upstream Py4J, and the other tests use in-process fake classes, so they validate the helper rather than Dataproc JavaMember. Linux/Databricks premerge will not catch a regression of the real call path. A test closer to Map.get().get(), plus evidence from a 2.2.86 image or an equivalent injected-Py4J IT run, would make this much stronger.
Signed-off-by: Allen Xu <allxu@nvidia.com>
wjxiz1992
changed the title
wjxiz1992
changed the title
wjxiz1992
changed the title
Greptile SummaryThis PR adds a process-local integration-test workaround that restores strong Py4J JavaMember container references when a runtime probe detects Dataproc’s weak-reference behavior.
Confidence Score: 4/5The PR appears safe to merge, with a non-blocking integration-test coverage issue around exercising the chained call through both CPU and GPU session configurations. The workaround is narrowly installed before the shared Spark session and is guarded by behavior detection and fail-open handling; the remaining concern is that its Spark-backed regression test uses only the default shared session path. Files Needing Attention: integration_tests/src/main/python/py4j_workaround_test.py Important Files Changed
Reviews (1): Last reviewed commit: "Address Py4J workaround review feedback" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
|
|
||
| def test_py4j_chained_scala_map_get(): | ||
| spark = get_spark_i_know_what_i_am_doing() | ||
| scala_map = spark.conf._jconf.getAll() | ||
| keys = scala_map.keys().iterator() | ||
|
|
||
| assert keys.hasNext() | ||
| key = keys.next() |
There was a problem hiding this comment.
This Spark-backed regression invokes the chained JVM map access only through the shared default session, so it can pass without confirming that the same call remains functional across the harness's CPU and GPU session configurations. Use the repository's CPU/GPU comparison or fallback assertion path for this integration test.
Rule Used: Integration tests must verify GPU execution using ... (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
There was a problem hiding this comment.
This regression targets driver-side Py4J JavaMember lifetime, not Spark query execution. with_cpu_session and with_gpu_session share the same SparkSession and gateway and only toggle RAPIDS SQL configuration, so they do not exercise a different lifetime path. The GPU/fallback assertion helpers require a DataFrame and execution plan, while this test intentionally exercises the exact scala_map.get(key).get() failure shape directly after the workaround is installed before SparkSession creation. Importing spark_session to use those helpers also runs the vulnerable _from_scala_map path during module initialization. A CPU/GPU split is therefore not applicable here.
There was a problem hiding this comment.
You're right — this regression is specifically about driver-side Py4J object lifetime, and the CPU/GPU helpers would not provide additional coverage because they reuse the same SparkSession and gateway while only changing RAPIDS SQL configuration. They also introduce the unrelated DataFrame/plan path and can trigger the vulnerable initialization earlier. The direct chained-call test is the appropriate coverage here, so the CPU/GPU session coverage request does not apply.
There was a problem hiding this comment.
🔵 Needs a closer look
It introduces a runtime monkey-patch of Py4J internals that could have subtle cross-test effects and warrants final human design review despite the added regression tests.
Pull request overview
This PR adds a process-local integration-test workaround for a Dataproc 2.2.86 Py4J behavior change where JavaMember.container becomes a weak reference, which can allow temporary JavaObject targets to be finalized before chained calls complete (e.g., scala_map.get(key).get()). The workaround probes the installed Py4J behavior and, only when needed, monkey-patches JavaMember.__init__ to restore a strong container reference for the test process prior to SparkSession creation.
Changes:
- Add a behavior-based probe and conditional monkey-patch to restore strong
JavaMember.containersemantics in the integration-test process. - Invoke the workaround during pytest session startup before creating the
SparkSession. - Add regression tests covering the chained-call shape, strong-reference retention, idempotency, upstream no-op behavior, and fail-open probe behavior.
File summaries
| File | Description |
|---|---|
| integration_tests/src/main/python/spark_init_internal.py | Adds conditional Py4J JavaMember container-lifetime workaround and applies it early in pytest session startup. |
| integration_tests/src/main/python/py4j_workaround_test.py | Adds regression tests validating detection/patching behavior and exercising the chained-call pattern that previously failed. |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 0
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
build |
wjxiz1992
changed the title
|
build |
JaCoCo production line coverage: not fully measurable locally — current nightly baseline artifacts were not applied; all JVM production modules are N/A for this Python-only integration-test harness change.
Refs #15805.
Description
Dataproc 2.2.86 changes Py4J's
JavaMember.containerfrom an upstream strong reference to a weak reference. A temporaryJavaObjectcan therefore be finalized, and its JVM target deleted, before a chained call such asscala_map.get(key).get()reachesJavaMember.__call__.This process-local workaround follows the runtime patch recommended by the Google Cloud Dataproc support team in Customer Care case
74970097(P2): restore the upstream strong-reference behavior while their internal investigation continues. The support case link is intentionally omitted; please contact me offline if the exact link is needed.The change applies only to the cudf-spark integration-test process before
SparkSessioncreation; it is not a Dataproc/Py4J platform fix and does not change the cudf-spark product JAR. It probes the installedJavaMemberbehavior with an inert container and installs the wrapper only when the container is actually stored as a weak reference. The wrapper is idempotent, and probe or initialization failures are logged without aborting the pytest worker. Standard upstream Py4J remains unmodified.Scope: this covers the Py4J missing-target failure family described in #15805. It does not address the separate ORC boolean-encoding signal reported on the same issue.
The regression tests cover the actual
scala_map.get(key).get()call shape used by the test harness, strong-reference retention after garbage collection, idempotent installation, the upstream no-op path, and fail-open behavior when the probe cannot run. This is a driver-side Py4J lifetime path rather than a Spark query: CPU and GPU session helpers share the same SparkSession and gateway and only toggle RAPIDS SQL configuration, so the focused test intentionally exercises the failing chained call directly instead of introducing an unrelated query or execution-plan assertion.Review focus: whether behavior-based detection is the right temporary boundary while Dataproc owns the permanent platform fix.
Validation completed locally:
python3 -m py_compilefor both changed Python files.4 passed, 1 deselected in 0.04s.5 passed in 25.74s.5 passed in 22.97s.git diff --check.AI assistance: The change and PR description were prepared with Codex assistance.
Checklists
Documentation
Testing
(Please provide the names of the existing tests in the PR description.)
Performance