[query] allow attaching debugger to Spark driver JVM

dev-docs/development-process.md

@@ -152,6 +152,14 @@ Run the Python query documentation tests:
 make -C hail doctest-query
 ```
 
+##### Attaching a debugger
+
+To help debug failing Python query tests, it is possible to attach the IntelliJ debugger to the JVM created by py4j/pyspark. The steps are:
+
+* In IntelliJ, make a "Remote JVM Debug" run/debug config, set to "listen to remote jvm", localhost, port 5005
+* run the debug config in IntelliJ
+* set `HAIL_DEBUG_JVM=1` when running hail python (e.g. `HAIL_DEBUG_JVM=1 make pytest PYTEST_ARGS='-k foo'`)
+
 #### Services
 
 Production, test and development versions of Hail Batch share one Kubernetes

hail/python/hail/backend/local_backend.py

@@ -30,7 +30,9 @@ def create(
     ) -> Py4JBackend:
         max_heap_size = jvm_heap_size or os.getenv('HAIL_LOCAL_BACKEND_HEAP_SIZE')
 
-        gateway = start_py4j_gateway(max_heap_size=max_heap_size)
+
+        attach_to_debugger = 'HAIL_DEBUG_JVM' in os.environ
+        gateway = start_py4j_gateway(max_heap_size=max_heap_size, attach_to_debugger=attach_to_debugger)
 
         try:
             raise_when_mismatched_hail_versions(gateway.jvm)

hail/python/hail/backend/py4j_backend.py

@@ -40,7 +40,7 @@
 _original = None
 
 
-def start_py4j_gateway(*, max_heap_size: str | None = None) -> JavaGateway:
+def start_py4j_gateway(*, max_heap_size: str | None = None, attach_to_debugger: bool = False) -> JavaGateway:
     spark_home = find_spark_home()
     info = local_jar_information()
     classpath = ':'.join([f'{spark_home}/jars/*', info.hail_jar, *info.extra_classpath])
@@ -49,6 +49,9 @@ def start_py4j_gateway(*, max_heap_size: str | None = None) -> JavaGateway:
     if max_heap_size is not None:
         jvm_opts.append(f'-Xmx{max_heap_size}')
 
+    if attach_to_debugger:
+        jvm_opts.append('-agentlib:jdwp=transport=dt_socket,server=n,suspend=y,address=localhost:5005')
+
     py4j_jars = glob.glob(f'{spark_home}/jars/py4j-*.jar')
     if len(py4j_jars) == 0:
         raise ValueError(f'No py4j JAR found in {spark_home}/jars')

hail/python/hail/backend/spark_backend.py

@@ -46,6 +46,9 @@ def _configure_spark_classpath(conf: SparkConf):
             conf, 'spark.extraListeners', _append_csv('sparkmonitor.listener.JupyterSparkMonitorListener')
         )
 
+    if os.environ.get('HAIL_DEBUG_JVM'):
+        conf.set('spark.driver.extraJavaOptions', '-agentlib:jdwp=transport=dt_socket,server=n,suspend=y,address=localhost:5005')
+
     _modify_spark_conf(conf, 'spark.jars', _append_csv(*classpath))
     if os.environ.get('AZURE_SPARK') == '1':
         print('AZURE_SPARK environment variable is set to "1", assuming you are in HDInsight.')