| # |
| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You 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 faulthandler |
| import gc |
| import os |
| import select |
| import signal |
| import socket |
| import sys |
| import time |
| import traceback |
| import uuid |
| from errno import EAGAIN, EINTR |
| from signal import SIG_DFL, SIG_IGN, SIGCHLD, SIGHUP, SIGINT, SIGTERM |
| from socket import AF_INET, AF_INET6, SOCK_STREAM, SOMAXCONN |
| from types import FrameType |
| from typing import Any, Optional |
| |
| from pyspark.errors import PySparkRuntimeError |
| from pyspark.serializers import UTF8Deserializer, read_int, write_int, write_with_length |
| from pyspark.util import enable_faulthandler |
| |
| |
| def compute_real_exit_code(exit_code: Any) -> int: |
| # SystemExit's code can be anything, but os._exit only accepts integer |
| if isinstance(exit_code, int): |
| return exit_code |
| else: |
| return 1 |
| |
| |
| def worker(sock: socket.socket, authenticated: bool) -> int: |
| """ |
| Called by a worker process after the fork(). |
| """ |
| signal.signal(SIGHUP, SIG_DFL) |
| signal.signal(SIGCHLD, SIG_DFL) |
| signal.signal(SIGTERM, SIG_DFL) |
| # restore the handler for SIGINT, |
| # it's useful for debugging (show the stacktrace before exit) |
| signal.signal(SIGINT, signal.default_int_handler) |
| |
| # Read the socket using fdopen instead of socket.makefile() because the latter |
| # seems to be very slow; note that we need to dup() the file descriptor because |
| # otherwise writes also cause a seek that makes us miss data on the read side. |
| buffer_size = int(os.environ.get("SPARK_BUFFER_SIZE", 65536)) |
| infile = os.fdopen(os.dup(sock.fileno()), "rb", buffer_size) |
| outfile = os.fdopen(os.dup(sock.fileno()), "wb", buffer_size) |
| |
| if not authenticated: |
| client_secret = UTF8Deserializer().loads(infile) |
| if os.environ["PYTHON_WORKER_FACTORY_SECRET"] == client_secret: |
| write_with_length("ok".encode("utf-8"), outfile) |
| outfile.flush() |
| else: |
| write_with_length("err".encode("utf-8"), outfile) |
| outfile.flush() |
| sock.close() |
| return 1 |
| |
| exit_code = 0 |
| |
| # We don't know what could happen when we import the worker module. We have to |
| # guarantee that no thread is spawned before we fork, so we have to import the |
| # worker module after fork. For example, both pandas and pyarrow starts some |
| # threads when they are imported. |
| if len(sys.argv) > 1 and sys.argv[1].startswith("pyspark"): |
| import importlib |
| |
| worker_module = importlib.import_module(sys.argv[1]) |
| worker_main = worker_module.main |
| else: |
| from pyspark.worker import main as worker_main |
| |
| try: |
| worker_main(infile, outfile) |
| except SystemExit as exc: |
| exit_code = compute_real_exit_code(exc.code) |
| finally: |
| try: |
| outfile.flush() |
| except Exception: |
| if os.environ.get("PYTHON_DAEMON_KILL_WORKER_ON_FLUSH_FAILURE", False): |
| faulthandler_log_path = os.environ.get("PYTHON_FAULTHANDLER_DIR", None) |
| if faulthandler_log_path: |
| faulthandler_log_path = os.path.join(faulthandler_log_path, str(os.getpid())) |
| with open( |
| faulthandler_log_path, "w", encoding="utf-8" |
| ) as faulthandler_log_file: |
| faulthandler.dump_traceback(file=faulthandler_log_file) |
| raise |
| else: |
| print( |
| "PySpark daemon failed to flush the output to the worker process:\n" |
| + traceback.format_exc(), |
| file=sys.stderr, |
| ) |
| return exit_code |
| |
| |
| def manager() -> None: |
| # Create a new process group to corral our children |
| os.setpgid(0, 0) |
| |
| is_unix_domain_sock = os.environ.get("PYTHON_UNIX_DOMAIN_ENABLED", "false").lower() == "true" |
| socket_path = None |
| |
| # Create a listening socket on the loopback interface |
| if is_unix_domain_sock: |
| assert "PYTHON_WORKER_FACTORY_SOCK_DIR" in os.environ |
| socket_path = os.path.join( |
| os.environ["PYTHON_WORKER_FACTORY_SOCK_DIR"], f".{uuid.uuid4()}.sock" |
| ) |
| listen_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) |
| listen_sock.bind(socket_path) |
| listen_sock.listen(max(1024, SOMAXCONN)) |
| listen_port = socket_path |
| elif os.environ.get("SPARK_PREFER_IPV6", "false").lower() == "true": |
| listen_sock = socket.socket(AF_INET6, SOCK_STREAM) |
| listen_sock.bind(("::1", 0, 0, 0)) |
| listen_sock.listen(max(1024, SOMAXCONN)) |
| listen_host, listen_port, _, _ = listen_sock.getsockname() |
| else: |
| listen_sock = socket.socket(AF_INET, SOCK_STREAM) |
| listen_sock.bind(("127.0.0.1", 0)) |
| listen_sock.listen(max(1024, SOMAXCONN)) |
| listen_host, listen_port = listen_sock.getsockname() |
| |
| # re-open stdin/stdout in 'wb' mode |
| stdin_bin = os.fdopen(sys.stdin.fileno(), "rb", 4) |
| stdout_bin = os.fdopen(sys.stdout.fileno(), "wb", 4) |
| if is_unix_domain_sock: |
| write_with_length(listen_port.encode("utf-8"), stdout_bin) |
| else: |
| write_int(listen_port, stdout_bin) |
| stdout_bin.flush() |
| |
| def shutdown(code: int) -> None: |
| if socket_path is not None and os.path.exists(socket_path): |
| os.remove(socket_path) |
| signal.signal(SIGTERM, SIG_DFL) |
| # Send SIGHUP to notify workers of shutdown |
| os.kill(0, SIGHUP) |
| sys.exit(code) |
| |
| def handle_sigterm(signal_number: int, frame: Optional[FrameType]) -> None: |
| shutdown(1) |
| |
| signal.signal(SIGTERM, handle_sigterm) # Gracefully exit on SIGTERM |
| signal.signal(SIGHUP, SIG_IGN) # Don't die on SIGHUP |
| signal.signal(SIGCHLD, SIG_IGN) |
| |
| reuse = os.environ.get("SPARK_REUSE_WORKER") |
| |
| # Initialization complete |
| try: |
| poller = None |
| if os.name == "posix": |
| # select.select has a known limit on the number of file descriptors |
| # it can handle. We use select.poll instead to avoid this limit. |
| poller = select.poll() |
| fd_reverse_map = {0: 0, listen_sock.fileno(): listen_sock} |
| poller.register(0, select.POLLIN) |
| poller.register(listen_sock, select.POLLIN) |
| |
| while True: |
| if poller is not None: |
| ready_fds = [] |
| # Unlike select, poll timeout is in millis. |
| for fd, event in poller.poll(1000): |
| if event & (select.POLLIN | select.POLLHUP): |
| # Data can be read (for POLLHUP peer hang up, so reads will return |
| # 0 bytes, in which case we want to break out - this is consistent |
| # with how select behaves). |
| ready_fds.append(fd_reverse_map[fd]) |
| else: |
| # Could be POLLERR or POLLNVAL (select would raise in this case). |
| raise PySparkRuntimeError(f"Polling error - event {event} on fd {fd}") |
| else: |
| # If poll is not available, use select. |
| ready_fds = select.select([0, listen_sock], [], [], 1)[0] |
| |
| if 0 in ready_fds: |
| try: |
| worker_pid = read_int(stdin_bin) |
| except EOFError: |
| # Spark told us to exit by closing stdin |
| shutdown(0) |
| try: |
| os.kill(worker_pid, signal.SIGKILL) |
| except OSError: |
| pass # process already died |
| |
| if listen_sock in ready_fds: |
| try: |
| sock, _ = listen_sock.accept() |
| except OSError as e: |
| if e.errno == EINTR: |
| continue |
| raise |
| |
| # Launch a worker process |
| try: |
| pid = os.fork() |
| except OSError as e: |
| if e.errno in (EAGAIN, EINTR): |
| time.sleep(1) |
| pid = os.fork() # error here will shutdown daemon |
| else: |
| outfile = sock.makefile(mode="wb") |
| write_int(e.errno, outfile) # Signal that the fork failed |
| outfile.flush() |
| outfile.close() |
| sock.close() |
| continue |
| |
| if pid == 0: |
| # in child process |
| with enable_faulthandler(): |
| if poller is not None: |
| poller.unregister(0) |
| poller.unregister(listen_sock) |
| listen_sock.close() |
| |
| # It should close the standard input in the child process so that |
| # Python native function executions stay intact. |
| # |
| # Note that if we just close the standard input (file descriptor 0), |
| # the lowest file descriptor (file descriptor 0) will be allocated, |
| # later when other file descriptors should happen to open. |
| # |
| # Therefore, here we redirects it to '/dev/null' by duplicating |
| # another file descriptor for '/dev/null' to the standard input (0). |
| # See SPARK-26175. |
| devnull = open(os.devnull, "r", encoding="utf-8") |
| os.dup2(devnull.fileno(), 0) |
| devnull.close() |
| |
| try: |
| # Acknowledge that the fork was successful |
| outfile = sock.makefile(mode="wb") |
| write_int(os.getpid(), outfile) |
| outfile.flush() |
| outfile.close() |
| authenticated = ( |
| os.environ.get("PYTHON_UNIX_DOMAIN_ENABLED", "false").lower() |
| == "true" |
| ) |
| while True: |
| code = worker(sock, authenticated) |
| if code == 0: |
| authenticated = True |
| if not reuse or code: |
| # wait for closing |
| try: |
| while sock.recv(1024): |
| pass |
| except Exception: |
| pass |
| break |
| gc.collect() |
| except BaseException: |
| traceback.print_exc() |
| os._exit(1) |
| else: |
| os._exit(0) |
| else: |
| sock.close() |
| |
| finally: |
| if poller is not None: |
| poller.unregister(0) |
| poller.unregister(listen_sock) |
| shutdown(1) |
| |
| |
| if __name__ == "__main__": |
| manager() |
