Work around Dataproc Py4J weak container references [reduced-it] [databricks]

integration_tests/src/main/python/py4j_workaround_test.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import gc
+import weakref
+
+from spark_init_internal import (
+    _apply_py4j_strong_ref_workaround_if_needed,
+    get_spark_i_know_what_i_am_doing,
+)
+
+
+class _Container:
+    pass
+
+
+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
+
+
+def test_py4j_weak_container_is_replaced_with_strong_reference():
+    class WeakJavaMember:
+        def __init__(self, name, container, target_id, gateway_client):
+            self.container = weakref.ref(container)
+
+    assert _apply_py4j_strong_ref_workaround_if_needed(WeakJavaMember)
+
+    container = _Container()
+    container_ref = weakref.ref(container)
+    member = WeakJavaMember('get', container, 'o1', None)
+    del container
+    gc.collect()
+
+    assert container_ref() is member.container
+
+
+def test_py4j_workaround_is_idempotent():
+    class WeakJavaMember:
+        def __init__(self, name, container, target_id, gateway_client):
+            self.container = weakref.ref(container)
+
+    assert _apply_py4j_strong_ref_workaround_if_needed(WeakJavaMember)
+    patched_init = WeakJavaMember.__init__
+
+    assert not _apply_py4j_strong_ref_workaround_if_needed(WeakJavaMember)
+    assert WeakJavaMember.__init__ is patched_init
+
+
+def test_py4j_upstream_strong_container_is_unchanged():
+    class StrongJavaMember:
+        def __init__(self, name, container, target_id, gateway_client):
+            self.container = container
+
+    original_init = StrongJavaMember.__init__
+
+    assert not _apply_py4j_strong_ref_workaround_if_needed(StrongJavaMember)
+    assert StrongJavaMember.__init__ is original_init

integration_tests/src/main/python/spark_init_internal.py

@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
 import logging
 import os
 import pytest
 import re
 import stat
 import traceback
+import weakref
 
 logging.basicConfig(
     format="%(asctime)s %(levelname)-8s %(message)s",
@@ -55,6 +57,52 @@ def conf_for_env(env_name):
 # does not raise NameError if session startup fails before assignment.
 _spark = None
 
+_PY4J_STRONG_REF_WORKAROUND_MARKER = '_spark_rapids_strong_ref_workaround'
+
+
+def _java_member_uses_weak_container(java_member_class):
+    class _ProbeContainer:
+        pass
+
+    class _ProbeGatewayProperty:
+        pool = None
+
+    class _ProbeGatewayClient:
+        gateway_property = _ProbeGatewayProperty()
+        converters = ()
+
+    container = _ProbeContainer()
+    member = java_member_class(
+        '_spark_rapids_container_lifetime_probe', container, 'o0', _ProbeGatewayClient())
+    return isinstance(member.container, weakref.ReferenceType)
+
+
+def _apply_py4j_strong_ref_workaround_if_needed(java_member_class=None):
+    """Restore upstream JavaMember container lifetime semantics when needed."""
+    if java_member_class is None:
+        from py4j.java_gateway import JavaMember
+        java_member_class = JavaMember
+
+    original_init = java_member_class.__init__
+    if getattr(original_init, _PY4J_STRONG_REF_WORKAROUND_MARKER, False):
+        return False
+    if not _java_member_uses_weak_container(java_member_class):
+        return False
+
+    @functools.wraps(original_init)
+    def strong_ref_init(self, name, container, target_id, gateway_client):
+        original_init(self, name, container, target_id, gateway_client)
+        if isinstance(self.container, weakref.ReferenceType):
+            self.container = container
+
+    setattr(strong_ref_init, _PY4J_STRONG_REF_WORKAROUND_MARKER, True)
+    java_member_class.__init__ = strong_ref_init
+    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 findspark_init():
     import findspark
     findspark.init()
@@ -120,6 +168,10 @@ def pytest_sessionstart(session):
     import pyspark
     from py4j.java_gateway import java_import
 
+    # Dataproc 2.2.86 can delete temporary JVM targets before chained Py4J calls complete.
+    # Keep this before SparkSession creation. The behavior probe makes it a no-op for upstream Py4J.
+    _apply_py4j_strong_ref_workaround_if_needed()
+
     # Force the RapidsPlugin to be enabled, so it blows up if the classpath is not set properly
     # DO NOT SET ANY OTHER CONFIGS HERE!!!
     # due to bugs in pyspark/pytest it looks like any configs set here