[SPARK-54441][INFRA] Enable coverage data for workers

python/.coveragerc

@@ -35,3 +35,8 @@ omit =
     */pyspark/streaming/tests/*
     */pyspark/tests/*
     */pyspark/testing/tests/*
+
+[paths]
+source =
+    pyspark/
+    lib/pyspark.zip/pyspark

python/test_coverage/sitecustomize.py

@@ -21,6 +21,32 @@
 # variable is set or not. If set, it starts to run the coverage.
 try:
     import coverage
-    coverage.process_startup()
+    cov = coverage.process_startup()
+    if cov:
+        import os
+
+        def patch_worker():
+            # If it's a worker forked from the daemon, we need to patch it to save
+            # the coverage data. Otherwise the worker will be killed by a signal and
+            # the coverage data will not be saved.
+            import sys
+            frame = sys._getframe(1)
+            if (
+                frame.f_code.co_name == "manager" and
+                "daemon.py" in frame.f_code.co_filename and
+                "worker" in frame.f_globals
+            ):
+
+                def save_when_exit(func):
+                    def wrapper(*args, **kwargs):
+                        result = func(*args, **kwargs)
+                        cov.save()
+                        return result
+                    return wrapper
+
+                frame.f_globals["worker"] = save_when_exit(frame.f_globals["worker"])
+
+        os.register_at_fork(after_in_child=patch_worker)
+
 except ImportError:
     pass