You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ue...@apache.org on 2025/11/18 21:52:25 UTC
(spark) branch master updated: [SPARK-54349][PYTHON] Refactor code a bit to simplify faulthandler integration extension
This is an automated email from the ASF dual-hosted git repository.
ueshin pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 6227fbab0308 [SPARK-54349][PYTHON] Refactor code a bit to simplify faulthandler integration extension
6227fbab0308 is described below
commit 6227fbab03082a19ad04ab300b89a721334eb9bb
Author: antban <dm...@gmail.com>
AuthorDate: Tue Nov 18 13:52:13 2025 -0800
[SPARK-54349][PYTHON] Refactor code a bit to simplify faulthandler integration extension
### What changes were proposed in this pull request?
There are many places where pyspark is trying to integrate with faulthandler and use the same functionality to dump stack traces/record thread dumps.
In order to reduce the complexity of the integration and ease the extension of integration it makes sense to technically refactor the code to use the same code in all the places.
### Why are the changes needed?
Improves developer experience.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
By existing unit tests
### Was this patch authored or co-authored using generative AI tooling?
Closes #53016 from antban/simplify-faulthandler-integration.
Authored-by: antban <dm...@gmail.com>
Signed-off-by: Takuya Ueshin <ue...@databricks.com>
---
python/pyspark/sql/worker/analyze_udtf.py | 27 +++------
.../pyspark/sql/worker/commit_data_source_write.py | 27 +++------
python/pyspark/sql/worker/create_data_source.py | 27 +++------
.../sql/worker/data_source_pushdown_filters.py | 27 +++------
python/pyspark/sql/worker/lookup_data_sources.py | 27 +++------
python/pyspark/sql/worker/plan_data_source_read.py | 27 +++------
.../sql/worker/python_streaming_sink_runner.py | 27 +++------
.../pyspark/sql/worker/write_into_data_source.py | 27 +++------
python/pyspark/util.py | 66 ++++++++++++++++++++++
python/pyspark/worker.py | 29 +++-------
10 files changed, 138 insertions(+), 173 deletions(-)
diff --git a/python/pyspark/sql/worker/analyze_udtf.py b/python/pyspark/sql/worker/analyze_udtf.py
index 665b1297fbc1..526cb316862c 100644
--- a/python/pyspark/sql/worker/analyze_udtf.py
+++ b/python/pyspark/sql/worker/analyze_udtf.py
@@ -15,7 +15,6 @@
# limitations under the License.
#
-import faulthandler
import inspect
import os
import sys
@@ -35,7 +34,12 @@ from pyspark.serializers import (
from pyspark.sql.functions import OrderingColumn, PartitioningColumn, SelectedColumn
from pyspark.sql.types import _parse_datatype_json_string, StructType
from pyspark.sql.udtf import AnalyzeArgument, AnalyzeResult
-from pyspark.util import handle_worker_exception, local_connect_and_auth
+from pyspark.util import (
+ handle_worker_exception,
+ local_connect_and_auth,
+ with_faulthandler,
+ start_faulthandler_periodic_traceback,
+)
from pyspark.worker_util import (
check_python_version,
read_command,
@@ -100,6 +104,7 @@ def read_arguments(infile: IO) -> Tuple[List[AnalyzeArgument], Dict[str, Analyze
return args, kwargs
+@with_faulthandler
def main(infile: IO, outfile: IO) -> None:
"""
Runs the Python UDTF's `analyze` static method.
@@ -108,18 +113,10 @@ def main(infile: IO, outfile: IO) -> None:
in JVM and receive the Python UDTF and its arguments for the `analyze` static method,
and call the `analyze` static method, and send back a AnalyzeResult as a result of the method.
"""
- faulthandler_log_path = os.environ.get("PYTHON_FAULTHANDLER_DIR", None)
- tracebackDumpIntervalSeconds = os.environ.get("PYTHON_TRACEBACK_DUMP_INTERVAL_SECONDS", None)
try:
- if faulthandler_log_path:
- faulthandler_log_path = os.path.join(faulthandler_log_path, str(os.getpid()))
- faulthandler_log_file = open(faulthandler_log_path, "w")
- faulthandler.enable(file=faulthandler_log_file)
-
check_python_version(infile)
- if tracebackDumpIntervalSeconds is not None and int(tracebackDumpIntervalSeconds) > 0:
- faulthandler.dump_traceback_later(int(tracebackDumpIntervalSeconds), repeat=True)
+ start_faulthandler_periodic_traceback()
memory_limit_mb = int(os.environ.get("PYSPARK_PLANNER_MEMORY_MB", "-1"))
setup_memory_limits(memory_limit_mb)
@@ -266,11 +263,6 @@ def main(infile: IO, outfile: IO) -> None:
except BaseException as e:
handle_worker_exception(e, outfile)
sys.exit(-1)
- finally:
- if faulthandler_log_path:
- faulthandler.disable()
- faulthandler_log_file.close()
- os.remove(faulthandler_log_path)
send_accumulator_updates(outfile)
@@ -282,9 +274,6 @@ def main(infile: IO, outfile: IO) -> None:
write_int(SpecialLengths.END_OF_DATA_SECTION, outfile)
sys.exit(-1)
- # Force to cancel dump_traceback_later
- faulthandler.cancel_dump_traceback_later()
-
if __name__ == "__main__":
# Read information about how to connect back to the JVM from the environment.
diff --git a/python/pyspark/sql/worker/commit_data_source_write.py b/python/pyspark/sql/worker/commit_data_source_write.py
index fb82b65f3122..37fee6ad8357 100644
--- a/python/pyspark/sql/worker/commit_data_source_write.py
+++ b/python/pyspark/sql/worker/commit_data_source_write.py
@@ -14,7 +14,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
-import faulthandler
import os
import sys
from typing import IO
@@ -29,7 +28,12 @@ from pyspark.serializers import (
SpecialLengths,
)
from pyspark.sql.datasource import DataSourceWriter, WriterCommitMessage
-from pyspark.util import handle_worker_exception, local_connect_and_auth
+from pyspark.util import (
+ handle_worker_exception,
+ local_connect_and_auth,
+ with_faulthandler,
+ start_faulthandler_periodic_traceback,
+)
from pyspark.worker_util import (
check_python_version,
pickleSer,
@@ -40,6 +44,7 @@ from pyspark.worker_util import (
)
+@with_faulthandler
def main(infile: IO, outfile: IO) -> None:
"""
Main method for committing or aborting a data source write operation.
@@ -49,18 +54,10 @@ def main(infile: IO, outfile: IO) -> None:
responsible for invoking either the `commit` or the `abort` method on a data source
writer instance, given a list of commit messages.
"""
- faulthandler_log_path = os.environ.get("PYTHON_FAULTHANDLER_DIR", None)
- tracebackDumpIntervalSeconds = os.environ.get("PYTHON_TRACEBACK_DUMP_INTERVAL_SECONDS", None)
try:
- if faulthandler_log_path:
- faulthandler_log_path = os.path.join(faulthandler_log_path, str(os.getpid()))
- faulthandler_log_file = open(faulthandler_log_path, "w")
- faulthandler.enable(file=faulthandler_log_file)
-
check_python_version(infile)
- if tracebackDumpIntervalSeconds is not None and int(tracebackDumpIntervalSeconds) > 0:
- faulthandler.dump_traceback_later(int(tracebackDumpIntervalSeconds), repeat=True)
+ start_faulthandler_periodic_traceback()
memory_limit_mb = int(os.environ.get("PYSPARK_PLANNER_MEMORY_MB", "-1"))
setup_memory_limits(memory_limit_mb)
@@ -106,11 +103,6 @@ def main(infile: IO, outfile: IO) -> None:
except BaseException as e:
handle_worker_exception(e, outfile)
sys.exit(-1)
- finally:
- if faulthandler_log_path:
- faulthandler.disable()
- faulthandler_log_file.close()
- os.remove(faulthandler_log_path)
send_accumulator_updates(outfile)
@@ -122,9 +114,6 @@ def main(infile: IO, outfile: IO) -> None:
write_int(SpecialLengths.END_OF_DATA_SECTION, outfile)
sys.exit(-1)
- # Force to cancel dump_traceback_later
- faulthandler.cancel_dump_traceback_later()
-
if __name__ == "__main__":
# Read information about how to connect back to the JVM from the environment.
diff --git a/python/pyspark/sql/worker/create_data_source.py b/python/pyspark/sql/worker/create_data_source.py
index 15e8fdc618e2..bf6ceda41ffb 100644
--- a/python/pyspark/sql/worker/create_data_source.py
+++ b/python/pyspark/sql/worker/create_data_source.py
@@ -14,7 +14,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
-import faulthandler
import inspect
import os
import sys
@@ -32,7 +31,12 @@ from pyspark.serializers import (
)
from pyspark.sql.datasource import DataSource, CaseInsensitiveDict
from pyspark.sql.types import _parse_datatype_json_string, StructType
-from pyspark.util import handle_worker_exception, local_connect_and_auth
+from pyspark.util import (
+ handle_worker_exception,
+ local_connect_and_auth,
+ with_faulthandler,
+ start_faulthandler_periodic_traceback,
+)
from pyspark.worker_util import (
check_python_version,
read_command,
@@ -45,6 +49,7 @@ from pyspark.worker_util import (
)
+@with_faulthandler
def main(infile: IO, outfile: IO) -> None:
"""
Main method for creating a Python data source instance.
@@ -62,18 +67,10 @@ def main(infile: IO, outfile: IO) -> None:
This process then creates a `DataSource` instance using the above information and
sends the pickled instance as well as the schema back to the JVM.
"""
- faulthandler_log_path = os.environ.get("PYTHON_FAULTHANDLER_DIR", None)
- tracebackDumpIntervalSeconds = os.environ.get("PYTHON_TRACEBACK_DUMP_INTERVAL_SECONDS", None)
try:
- if faulthandler_log_path:
- faulthandler_log_path = os.path.join(faulthandler_log_path, str(os.getpid()))
- faulthandler_log_file = open(faulthandler_log_path, "w")
- faulthandler.enable(file=faulthandler_log_file)
-
check_python_version(infile)
- if tracebackDumpIntervalSeconds is not None and int(tracebackDumpIntervalSeconds) > 0:
- faulthandler.dump_traceback_later(int(tracebackDumpIntervalSeconds), repeat=True)
+ start_faulthandler_periodic_traceback()
memory_limit_mb = int(os.environ.get("PYSPARK_PLANNER_MEMORY_MB", "-1"))
setup_memory_limits(memory_limit_mb)
@@ -172,11 +169,6 @@ def main(infile: IO, outfile: IO) -> None:
except BaseException as e:
handle_worker_exception(e, outfile)
sys.exit(-1)
- finally:
- if faulthandler_log_path:
- faulthandler.disable()
- faulthandler_log_file.close()
- os.remove(faulthandler_log_path)
send_accumulator_updates(outfile)
@@ -188,9 +180,6 @@ def main(infile: IO, outfile: IO) -> None:
write_int(SpecialLengths.END_OF_DATA_SECTION, outfile)
sys.exit(-1)
- # Force to cancel dump_traceback_later
- faulthandler.cancel_dump_traceback_later()
-
if __name__ == "__main__":
# Read information about how to connect back to the JVM from the environment.
diff --git a/python/pyspark/sql/worker/data_source_pushdown_filters.py b/python/pyspark/sql/worker/data_source_pushdown_filters.py
index 8601521bcfb1..7d255e1dbf77 100644
--- a/python/pyspark/sql/worker/data_source_pushdown_filters.py
+++ b/python/pyspark/sql/worker/data_source_pushdown_filters.py
@@ -16,7 +16,6 @@
#
import base64
-import faulthandler
import json
import os
import sys
@@ -49,7 +48,12 @@ from pyspark.sql.datasource import (
)
from pyspark.sql.types import StructType, VariantVal, _parse_datatype_json_string
from pyspark.sql.worker.plan_data_source_read import write_read_func_and_partitions
-from pyspark.util import handle_worker_exception, local_connect_and_auth
+from pyspark.util import (
+ handle_worker_exception,
+ local_connect_and_auth,
+ with_faulthandler,
+ start_faulthandler_periodic_traceback,
+)
from pyspark.worker_util import (
check_python_version,
pickleSer,
@@ -119,6 +123,7 @@ def deserializeFilter(jsonDict: dict) -> Filter:
return filter
+@with_faulthandler
def main(infile: IO, outfile: IO) -> None:
"""
Main method for planning a data source read with filter pushdown.
@@ -140,18 +145,10 @@ def main(infile: IO, outfile: IO) -> None:
on the reader and determines which filters are supported. The indices of the supported
filters are sent back to the JVM, along with the list of partitions and the read function.
"""
- faulthandler_log_path = os.environ.get("PYTHON_FAULTHANDLER_DIR", None)
- tracebackDumpIntervalSeconds = os.environ.get("PYTHON_TRACEBACK_DUMP_INTERVAL_SECONDS", None)
try:
- if faulthandler_log_path:
- faulthandler_log_path = os.path.join(faulthandler_log_path, str(os.getpid()))
- faulthandler_log_file = open(faulthandler_log_path, "w")
- faulthandler.enable(file=faulthandler_log_file)
-
check_python_version(infile)
- if tracebackDumpIntervalSeconds is not None and int(tracebackDumpIntervalSeconds) > 0:
- faulthandler.dump_traceback_later(int(tracebackDumpIntervalSeconds), repeat=True)
+ start_faulthandler_periodic_traceback()
memory_limit_mb = int(os.environ.get("PYSPARK_PLANNER_MEMORY_MB", "-1"))
setup_memory_limits(memory_limit_mb)
@@ -258,11 +255,6 @@ def main(infile: IO, outfile: IO) -> None:
except BaseException as e:
handle_worker_exception(e, outfile)
sys.exit(-1)
- finally:
- if faulthandler_log_path:
- faulthandler.disable()
- faulthandler_log_file.close()
- os.remove(faulthandler_log_path)
send_accumulator_updates(outfile)
@@ -274,9 +266,6 @@ def main(infile: IO, outfile: IO) -> None:
write_int(SpecialLengths.END_OF_DATA_SECTION, outfile)
sys.exit(-1)
- # Force to cancel dump_traceback_later
- faulthandler.cancel_dump_traceback_later()
-
if __name__ == "__main__":
# Read information about how to connect back to the JVM from the environment.
diff --git a/python/pyspark/sql/worker/lookup_data_sources.py b/python/pyspark/sql/worker/lookup_data_sources.py
index eeb84263d445..b23903cac8cb 100644
--- a/python/pyspark/sql/worker/lookup_data_sources.py
+++ b/python/pyspark/sql/worker/lookup_data_sources.py
@@ -14,7 +14,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
-import faulthandler
from importlib import import_module
from pkgutil import iter_modules
import os
@@ -29,7 +28,12 @@ from pyspark.serializers import (
SpecialLengths,
)
from pyspark.sql.datasource import DataSource
-from pyspark.util import handle_worker_exception, local_connect_and_auth
+from pyspark.util import (
+ handle_worker_exception,
+ local_connect_and_auth,
+ with_faulthandler,
+ start_faulthandler_periodic_traceback,
+)
from pyspark.worker_util import (
check_python_version,
pickleSer,
@@ -40,6 +44,7 @@ from pyspark.worker_util import (
)
+@with_faulthandler
def main(infile: IO, outfile: IO) -> None:
"""
Main method for looking up the available Python Data Sources in Python path.
@@ -51,18 +56,10 @@ def main(infile: IO, outfile: IO) -> None:
This is responsible for searching the available Python Data Sources so they can be
statically registered automatically.
"""
- faulthandler_log_path = os.environ.get("PYTHON_FAULTHANDLER_DIR", None)
- tracebackDumpIntervalSeconds = os.environ.get("PYTHON_TRACEBACK_DUMP_INTERVAL_SECONDS", None)
try:
- if faulthandler_log_path:
- faulthandler_log_path = os.path.join(faulthandler_log_path, str(os.getpid()))
- faulthandler_log_file = open(faulthandler_log_path, "w")
- faulthandler.enable(file=faulthandler_log_file)
-
check_python_version(infile)
- if tracebackDumpIntervalSeconds is not None and int(tracebackDumpIntervalSeconds) > 0:
- faulthandler.dump_traceback_later(int(tracebackDumpIntervalSeconds), repeat=True)
+ start_faulthandler_periodic_traceback()
memory_limit_mb = int(os.environ.get("PYSPARK_PLANNER_MEMORY_MB", "-1"))
setup_memory_limits(memory_limit_mb)
@@ -89,11 +86,6 @@ def main(infile: IO, outfile: IO) -> None:
except BaseException as e:
handle_worker_exception(e, outfile)
sys.exit(-1)
- finally:
- if faulthandler_log_path:
- faulthandler.disable()
- faulthandler_log_file.close()
- os.remove(faulthandler_log_path)
send_accumulator_updates(outfile)
@@ -105,9 +97,6 @@ def main(infile: IO, outfile: IO) -> None:
write_int(SpecialLengths.END_OF_DATA_SECTION, outfile)
sys.exit(-1)
- # Force to cancel dump_traceback_later
- faulthandler.cancel_dump_traceback_later()
-
if __name__ == "__main__":
# Read information about how to connect back to the JVM from the environment.
diff --git a/python/pyspark/sql/worker/plan_data_source_read.py b/python/pyspark/sql/worker/plan_data_source_read.py
index db79e58f2ec4..51036f17586f 100644
--- a/python/pyspark/sql/worker/plan_data_source_read.py
+++ b/python/pyspark/sql/worker/plan_data_source_read.py
@@ -15,7 +15,6 @@
# limitations under the License.
#
-import faulthandler
import os
import sys
import functools
@@ -47,7 +46,12 @@ from pyspark.sql.types import (
BinaryType,
StructType,
)
-from pyspark.util import handle_worker_exception, local_connect_and_auth
+from pyspark.util import (
+ handle_worker_exception,
+ local_connect_and_auth,
+ with_faulthandler,
+ start_faulthandler_periodic_traceback,
+)
from pyspark.worker_util import (
check_python_version,
read_command,
@@ -267,6 +271,7 @@ def write_read_func_and_partitions(
write_int(0, outfile)
+@with_faulthandler
def main(infile: IO, outfile: IO) -> None:
"""
Main method for planning a data source read.
@@ -287,18 +292,10 @@ def main(infile: IO, outfile: IO) -> None:
The partition values and the Arrow Batch are then serialized and sent back to the JVM
via the socket.
"""
- faulthandler_log_path = os.environ.get("PYTHON_FAULTHANDLER_DIR", None)
- tracebackDumpIntervalSeconds = os.environ.get("PYTHON_TRACEBACK_DUMP_INTERVAL_SECONDS", None)
try:
- if faulthandler_log_path:
- faulthandler_log_path = os.path.join(faulthandler_log_path, str(os.getpid()))
- faulthandler_log_file = open(faulthandler_log_path, "w")
- faulthandler.enable(file=faulthandler_log_file)
-
check_python_version(infile)
- if tracebackDumpIntervalSeconds is not None and int(tracebackDumpIntervalSeconds) > 0:
- faulthandler.dump_traceback_later(int(tracebackDumpIntervalSeconds), repeat=True)
+ start_faulthandler_periodic_traceback()
memory_limit_mb = int(os.environ.get("PYSPARK_PLANNER_MEMORY_MB", "-1"))
setup_memory_limits(memory_limit_mb)
@@ -402,11 +399,6 @@ def main(infile: IO, outfile: IO) -> None:
except BaseException as e:
handle_worker_exception(e, outfile)
sys.exit(-1)
- finally:
- if faulthandler_log_path:
- faulthandler.disable()
- faulthandler_log_file.close()
- os.remove(faulthandler_log_path)
send_accumulator_updates(outfile)
@@ -418,9 +410,6 @@ def main(infile: IO, outfile: IO) -> None:
write_int(SpecialLengths.END_OF_DATA_SECTION, outfile)
sys.exit(-1)
- # Force to cancel dump_traceback_later
- faulthandler.cancel_dump_traceback_later()
-
if __name__ == "__main__":
# Read information about how to connect back to the JVM from the environment.
diff --git a/python/pyspark/sql/worker/python_streaming_sink_runner.py b/python/pyspark/sql/worker/python_streaming_sink_runner.py
index ed6907ce5b63..5ca3307fca33 100644
--- a/python/pyspark/sql/worker/python_streaming_sink_runner.py
+++ b/python/pyspark/sql/worker/python_streaming_sink_runner.py
@@ -15,7 +15,6 @@
# limitations under the License.
#
-import faulthandler
import os
import sys
from typing import IO
@@ -35,7 +34,12 @@ from pyspark.sql.types import (
_parse_datatype_json_string,
StructType,
)
-from pyspark.util import handle_worker_exception, local_connect_and_auth
+from pyspark.util import (
+ handle_worker_exception,
+ local_connect_and_auth,
+ with_faulthandler,
+ start_faulthandler_periodic_traceback,
+)
from pyspark.worker_util import (
check_python_version,
read_command,
@@ -48,6 +52,7 @@ from pyspark.worker_util import (
)
+@with_faulthandler
def main(infile: IO, outfile: IO) -> None:
"""
Main method for committing or aborting a data source streaming write operation.
@@ -57,18 +62,10 @@ def main(infile: IO, outfile: IO) -> None:
responsible for invoking either the `commit` or the `abort` method on a data source
writer instance, given a list of commit messages.
"""
- faulthandler_log_path = os.environ.get("PYTHON_FAULTHANDLER_DIR", None)
- tracebackDumpIntervalSeconds = os.environ.get("PYTHON_TRACEBACK_DUMP_INTERVAL_SECONDS", None)
try:
- if faulthandler_log_path:
- faulthandler_log_path = os.path.join(faulthandler_log_path, str(os.getpid()))
- faulthandler_log_file = open(faulthandler_log_path, "w")
- faulthandler.enable(file=faulthandler_log_file)
-
check_python_version(infile)
- if tracebackDumpIntervalSeconds is not None and int(tracebackDumpIntervalSeconds) > 0:
- faulthandler.dump_traceback_later(int(tracebackDumpIntervalSeconds), repeat=True)
+ start_faulthandler_periodic_traceback()
setup_spark_files(infile)
setup_broadcasts(infile)
@@ -138,11 +135,6 @@ def main(infile: IO, outfile: IO) -> None:
except BaseException as e:
handle_worker_exception(e, outfile)
sys.exit(-1)
- finally:
- if faulthandler_log_path:
- faulthandler.disable()
- faulthandler_log_file.close()
- os.remove(faulthandler_log_path)
send_accumulator_updates(outfile)
@@ -154,9 +146,6 @@ def main(infile: IO, outfile: IO) -> None:
write_int(SpecialLengths.END_OF_DATA_SECTION, outfile)
sys.exit(-1)
- # Force to cancel dump_traceback_later
- faulthandler.cancel_dump_traceback_later()
-
if __name__ == "__main__":
# Read information about how to connect back to the JVM from the environment.
diff --git a/python/pyspark/sql/worker/write_into_data_source.py b/python/pyspark/sql/worker/write_into_data_source.py
index 917d0ca8e007..b8a54f8397dc 100644
--- a/python/pyspark/sql/worker/write_into_data_source.py
+++ b/python/pyspark/sql/worker/write_into_data_source.py
@@ -14,7 +14,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
-import faulthandler
import inspect
import os
import sys
@@ -46,7 +45,12 @@ from pyspark.sql.types import (
BinaryType,
_create_row,
)
-from pyspark.util import handle_worker_exception, local_connect_and_auth
+from pyspark.util import (
+ handle_worker_exception,
+ local_connect_and_auth,
+ with_faulthandler,
+ start_faulthandler_periodic_traceback,
+)
from pyspark.worker_util import (
check_python_version,
read_command,
@@ -59,6 +63,7 @@ from pyspark.worker_util import (
)
+@with_faulthandler
def main(infile: IO, outfile: IO) -> None:
"""
Main method for saving into a Python data source.
@@ -78,18 +83,10 @@ def main(infile: IO, outfile: IO) -> None:
instance and send a function using the writer instance that can be used
in mapInPandas/mapInArrow back to the JVM.
"""
- faulthandler_log_path = os.environ.get("PYTHON_FAULTHANDLER_DIR", None)
- tracebackDumpIntervalSeconds = os.environ.get("PYTHON_TRACEBACK_DUMP_INTERVAL_SECONDS", None)
try:
- if faulthandler_log_path:
- faulthandler_log_path = os.path.join(faulthandler_log_path, str(os.getpid()))
- faulthandler_log_file = open(faulthandler_log_path, "w")
- faulthandler.enable(file=faulthandler_log_file)
-
check_python_version(infile)
- if tracebackDumpIntervalSeconds is not None and int(tracebackDumpIntervalSeconds) > 0:
- faulthandler.dump_traceback_later(int(tracebackDumpIntervalSeconds), repeat=True)
+ start_faulthandler_periodic_traceback()
memory_limit_mb = int(os.environ.get("PYSPARK_PLANNER_MEMORY_MB", "-1"))
setup_memory_limits(memory_limit_mb)
@@ -264,11 +261,6 @@ def main(infile: IO, outfile: IO) -> None:
except BaseException as e:
handle_worker_exception(e, outfile)
sys.exit(-1)
- finally:
- if faulthandler_log_path:
- faulthandler.disable()
- faulthandler_log_file.close()
- os.remove(faulthandler_log_path)
send_accumulator_updates(outfile)
@@ -280,9 +272,6 @@ def main(infile: IO, outfile: IO) -> None:
write_int(SpecialLengths.END_OF_DATA_SECTION, outfile)
sys.exit(-1)
- # Force to cancel dump_traceback_later
- faulthandler.cancel_dump_traceback_later()
-
if __name__ == "__main__":
# Read information about how to connect back to the JVM from the environment.
diff --git a/python/pyspark/util.py b/python/pyspark/util.py
index 501d0afbb163..9935206df177 100644
--- a/python/pyspark/util.py
+++ b/python/pyspark/util.py
@@ -18,6 +18,7 @@
import copy
import functools
+import faulthandler
import itertools
import os
import platform
@@ -917,6 +918,71 @@ def default_api_mode() -> str:
return "classic"
+class _FaulthandlerHelper:
+ def __init__(self) -> None:
+ self._log_path: Optional[str] = None
+ self._log_file: Optional[TextIO] = None
+ self._periodic_traceback = False
+
+ def start(self) -> None:
+ if self._log_path:
+ raise Exception("Fault handler is already registered. No second registration allowed")
+ self._log_path = os.environ.get("PYTHON_FAULTHANDLER_DIR", None)
+ if self._log_path:
+ self._log_path = os.path.join(self._log_path, str(os.getpid()))
+ self._log_file = open(self._log_path, "w")
+
+ faulthandler.enable(file=self._log_file)
+
+ def stop(self) -> None:
+ if self._log_path:
+ faulthandler.disable()
+ if self._log_file:
+ self._log_file.close()
+ self._log_file = None
+ try:
+ os.remove(self._log_path)
+ finally:
+ self._log_path = None
+ if self._periodic_traceback:
+ faulthandler.cancel_dump_traceback_later()
+ self._periodic_traceback = False
+
+ def start_periodic_traceback(self) -> None:
+ # If the registration is already done - do nothing
+ if self._periodic_traceback:
+ return
+
+ traceback_dump_interval_seconds = os.environ.get(
+ "PYTHON_TRACEBACK_DUMP_INTERVAL_SECONDS", None
+ )
+ if traceback_dump_interval_seconds is not None and int(traceback_dump_interval_seconds) > 0:
+ self._periodic_traceback = True
+ faulthandler.dump_traceback_later(int(traceback_dump_interval_seconds), repeat=True)
+
+ def with_faulthandler(self, func: Callable) -> Callable:
+ """
+ Registers fault handler for the duration of function execution.
+ After function execution is over the faulthandler registration is cleaned as well,
+ including any files created for the integration.
+ """
+
+ @functools.wraps(func)
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
+ try:
+ self.start()
+ return func(*args, **kwargs)
+ finally:
+ self.stop()
+
+ return wrapper
+
+
+_faulthandler_helper = _FaulthandlerHelper()
+with_faulthandler = _faulthandler_helper.with_faulthandler
+start_faulthandler_periodic_traceback = _faulthandler_helper.start_periodic_traceback
+
+
if __name__ == "__main__":
if "pypy" not in platform.python_implementation().lower():
import doctest
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 6e34b041665a..4bae9f6dc48f 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -27,7 +27,6 @@ import inspect
import itertools
import json
from typing import Any, Callable, Iterable, Iterator, Optional, Tuple
-import faulthandler
from pyspark.accumulators import (
SpecialAccumulatorIds,
@@ -84,7 +83,12 @@ from pyspark.sql.types import (
_create_row,
_parse_datatype_json_string,
)
-from pyspark.util import fail_on_stopiteration, handle_worker_exception
+from pyspark.util import (
+ fail_on_stopiteration,
+ handle_worker_exception,
+ with_faulthandler,
+ start_faulthandler_periodic_traceback,
+)
from pyspark import shuffle
from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError
from pyspark.worker_util import (
@@ -3290,23 +3294,14 @@ def read_udfs(pickleSer, infile, eval_type):
return func, None, ser, ser
+@with_faulthandler
def main(infile, outfile):
- faulthandler_log_path = os.environ.get("PYTHON_FAULTHANDLER_DIR", None)
- tracebackDumpIntervalSeconds = os.environ.get("PYTHON_TRACEBACK_DUMP_INTERVAL_SECONDS", None)
try:
- if faulthandler_log_path:
- faulthandler_log_path = os.path.join(faulthandler_log_path, str(os.getpid()))
- faulthandler_log_file = open(faulthandler_log_path, "w")
- faulthandler.enable(file=faulthandler_log_file)
-
boot_time = time.time()
split_index = read_int(infile)
if split_index == -1: # for unit tests
sys.exit(-1)
-
- if tracebackDumpIntervalSeconds is not None and int(tracebackDumpIntervalSeconds) > 0:
- faulthandler.dump_traceback_later(int(tracebackDumpIntervalSeconds), repeat=True)
-
+ start_faulthandler_periodic_traceback()
check_python_version(infile)
# read inputs only for a barrier task
@@ -3397,11 +3392,6 @@ def main(infile, outfile):
except BaseException as e:
handle_worker_exception(e, outfile)
sys.exit(-1)
- finally:
- if faulthandler_log_path:
- faulthandler.disable()
- faulthandler_log_file.close()
- os.remove(faulthandler_log_path)
finish_time = time.time()
report_times(outfile, boot_time, init_time, finish_time)
write_long(shuffle.MemoryBytesSpilled, outfile)
@@ -3419,9 +3409,6 @@ def main(infile, outfile):
write_int(SpecialLengths.END_OF_DATA_SECTION, outfile)
sys.exit(-1)
- # Force to cancel dump_traceback_later
- faulthandler.cancel_dump_traceback_later()
-
if __name__ == "__main__":
# Read information about how to connect back to the JVM from the environment.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org