gaogaotiantian commented on code in PR #56907: URL: https://github.com/apache/spark/pull/56907#discussion_r3546611188
########## python/pyspark/sql/connect/local_server.py: ########## @@ -0,0 +1,212 @@ +# +# 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. +# + +""" +Long-lived local Spark Connect server for the opt-in reuse path. + +This module is launched as a *detached* child process by +``SparkSession._reuse_or_start_local_connect_server`` (in ``pyspark.sql.connect.session``) when +``SPARK_LOCAL_CONNECT_REUSE`` / ``spark.local.connect.reuse`` is enabled. It starts a regular +(classic) local Spark session with the Spark Connect plugin -- the same mechanism the in-process +``SparkSession._start_connect_server`` uses -- and then blocks, so one warm JVM and Spark Connect +server keeps serving many short-lived client processes. Each client connection gets its own +isolated server-side session, so session-local state (temp views, runtime SQL confs, isolated +artifacts) does not leak between runs. + +Once the server is accepting connections it writes a discovery file (host, the actually bound port, +the auth token, its pid and the Spark version) that later client processes read to reconnect. + +It is launched by file path rather than ``python -m`` so it does not require the Spark Connect Review Comment: This is feels wrong to me. A script in `sql/connect` that is executed by file - it does not fit the pattern. Why do we need to avoid requiring client requirements? This is already designed as a local debugging tool. I don't think we need to distinguish between "server" and "client". This script is executed by the user script which needs the client to work right? ########## python/pyspark/sql/connect/local_server.py: ########## @@ -0,0 +1,212 @@ +# +# 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. +# + +""" +Long-lived local Spark Connect server for the opt-in reuse path. + +This module is launched as a *detached* child process by +``SparkSession._reuse_or_start_local_connect_server`` (in ``pyspark.sql.connect.session``) when +``SPARK_LOCAL_CONNECT_REUSE`` / ``spark.local.connect.reuse`` is enabled. It starts a regular +(classic) local Spark session with the Spark Connect plugin -- the same mechanism the in-process +``SparkSession._start_connect_server`` uses -- and then blocks, so one warm JVM and Spark Connect +server keeps serving many short-lived client processes. Each client connection gets its own +isolated server-side session, so session-local state (temp views, runtime SQL confs, isolated +artifacts) does not leak between runs. + +Once the server is accepting connections it writes a discovery file (host, the actually bound port, +the auth token, its pid and the Spark version) that later client processes read to reconnect. + +It is launched by file path rather than ``python -m`` so it does not require the Spark Connect +*client* dependencies (grpc, etc.): a server only needs a classic PySpark install plus the Connect +server jar, like ``sbin/start-connect-server.sh``. It imports only the classic ``pyspark.sql`` API. +""" + +import sys + +# Launching by file path puts this file's directory -- pyspark/sql/connect -- at the front of +# sys.path, where modules such as `types` and `logging` shadow the standard library and break +# ordinary imports. Drop it before importing anything else (`import sys` cannot be shadowed); +# pyspark stays importable via the remaining sys.path entries. +if sys.path: + del sys.path[0] + +import argparse +import json +import os +import signal +import time +from typing import Any + + +def _write_discovery(path: str, host: str, port: int, token: str, version: str) -> None: + """Atomically write the discovery file with ``0600`` perms (it holds the auth token).""" + parent = os.path.dirname(path) + if parent and not os.path.isdir(parent): + os.makedirs(parent, exist_ok=True) + payload = { + "host": host, + "port": port, + "token": token, + "pid": os.getpid(), + "spark_version": version, + } + tmp = "{}.{}.tmp".format(path, os.getpid()) + fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as f: + f.write(json.dumps(payload)) + os.replace(tmp, path) + + +def _remove_discovery_if_ours(path: str) -> None: + """Remove the discovery file, but only if it still points at this process.""" + try: + with open(path, "r") as f: + disc = json.load(f) + except (OSError, ValueError): + return + if disc.get("pid") == os.getpid(): + try: + os.remove(path) + except OSError: + pass + + +def _has_active_sessions(spark: Any) -> bool: + """Best-effort check for whether any Spark Connect session is currently registered. + + Used only by the idle-shutdown reaper. Any failure (server not started yet, API drift, py4j + error) returns ``True`` so the reaper never terminates a server it cannot inspect. + """ + jvm = spark.sparkContext._jvm + service = getattr( + getattr(jvm, "org.apache.spark.sql.connect.service.SparkConnectService$"), + "MODULE$", + ) + if not service.started(): + return True + return not service.sessionManager().listActiveSessions().isEmpty() + + +def _bound_port(spark: Any, requested_port: int) -> int: + """Return the port the Connect server actually bound (``requested_port`` may have been 0).""" + jvm = spark.sparkContext._jvm + service = getattr( + getattr(jvm, "org.apache.spark.sql.connect.service.SparkConnectService$"), + "MODULE$", + ) + try: + return int(service.localPort()) + except Exception: + return requested_port + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--master", default="local[*]") + parser.add_argument("--port", type=int, default=15002) + parser.add_argument("--token", default=None) + parser.add_argument("--discovery", required=True) + parser.add_argument( + "--conf-file", + default=None, + help="path to a JSON file of extra SparkConf entries to seed the server with", + ) + parser.add_argument( + "--idle-timeout", + type=float, + default=3600.0, Review Comment: For local development, this seems a bit too long. ########## python/pyspark/sql/connect/local_server.py: ########## @@ -0,0 +1,212 @@ +# +# 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. +# + +""" +Long-lived local Spark Connect server for the opt-in reuse path. + +This module is launched as a *detached* child process by +``SparkSession._reuse_or_start_local_connect_server`` (in ``pyspark.sql.connect.session``) when +``SPARK_LOCAL_CONNECT_REUSE`` / ``spark.local.connect.reuse`` is enabled. It starts a regular +(classic) local Spark session with the Spark Connect plugin -- the same mechanism the in-process +``SparkSession._start_connect_server`` uses -- and then blocks, so one warm JVM and Spark Connect +server keeps serving many short-lived client processes. Each client connection gets its own +isolated server-side session, so session-local state (temp views, runtime SQL confs, isolated +artifacts) does not leak between runs. + +Once the server is accepting connections it writes a discovery file (host, the actually bound port, +the auth token, its pid and the Spark version) that later client processes read to reconnect. + +It is launched by file path rather than ``python -m`` so it does not require the Spark Connect +*client* dependencies (grpc, etc.): a server only needs a classic PySpark install plus the Connect +server jar, like ``sbin/start-connect-server.sh``. It imports only the classic ``pyspark.sql`` API. +""" + +import sys + +# Launching by file path puts this file's directory -- pyspark/sql/connect -- at the front of +# sys.path, where modules such as `types` and `logging` shadow the standard library and break +# ordinary imports. Drop it before importing anything else (`import sys` cannot be shadowed); +# pyspark stays importable via the remaining sys.path entries. +if sys.path: + del sys.path[0] + +import argparse +import json +import os +import signal +import time +from typing import Any + + +def _write_discovery(path: str, host: str, port: int, token: str, version: str) -> None: + """Atomically write the discovery file with ``0600`` perms (it holds the auth token).""" + parent = os.path.dirname(path) + if parent and not os.path.isdir(parent): + os.makedirs(parent, exist_ok=True) + payload = { + "host": host, + "port": port, + "token": token, + "pid": os.getpid(), + "spark_version": version, + } + tmp = "{}.{}.tmp".format(path, os.getpid()) + fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as f: + f.write(json.dumps(payload)) + os.replace(tmp, path) + + +def _remove_discovery_if_ours(path: str) -> None: + """Remove the discovery file, but only if it still points at this process.""" + try: + with open(path, "r") as f: + disc = json.load(f) + except (OSError, ValueError): + return + if disc.get("pid") == os.getpid(): + try: + os.remove(path) + except OSError: + pass + + +def _has_active_sessions(spark: Any) -> bool: + """Best-effort check for whether any Spark Connect session is currently registered. + + Used only by the idle-shutdown reaper. Any failure (server not started yet, API drift, py4j + error) returns ``True`` so the reaper never terminates a server it cannot inspect. + """ + jvm = spark.sparkContext._jvm + service = getattr( + getattr(jvm, "org.apache.spark.sql.connect.service.SparkConnectService$"), + "MODULE$", + ) + if not service.started(): + return True + return not service.sessionManager().listActiveSessions().isEmpty() + + +def _bound_port(spark: Any, requested_port: int) -> int: + """Return the port the Connect server actually bound (``requested_port`` may have been 0).""" + jvm = spark.sparkContext._jvm + service = getattr( + getattr(jvm, "org.apache.spark.sql.connect.service.SparkConnectService$"), + "MODULE$", + ) + try: + return int(service.localPort()) + except Exception: + return requested_port + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--master", default="local[*]") + parser.add_argument("--port", type=int, default=15002) + parser.add_argument("--token", default=None) + parser.add_argument("--discovery", required=True) + parser.add_argument( + "--conf-file", + default=None, + help="path to a JSON file of extra SparkConf entries to seed the server with", + ) + parser.add_argument( + "--idle-timeout", + type=float, + default=3600.0, + help="seconds with no active session after which the server self-terminates; <=0 disables", + ) + parser.add_argument("--poll-interval", type=float, default=60.0) + args = parser.parse_args() + + # Build a CLASSIC session: the connect-mode env vars would otherwise divert us into a client. + for var in ("SPARK_REMOTE", "SPARK_LOCAL_REMOTE", "SPARK_CONNECT_MODE_ENABLED"): + os.environ.pop(var, None) + if args.token: + os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = args.token Review Comment: What's the relationship between this env var and the config below? ########## python/pyspark/sql/connect/session.py: ########## @@ -1247,6 +1249,372 @@ def _start_connect_server(master: str, opts: Dict[str, Any]) -> None: messageParameters={}, ) + # Opt-in reuse of a persistent local Spark Connect server. By default ``.remote("local[*]")`` + # boots a fresh in-process server every process (see ``_start_connect_server`` above); when + # reuse is enabled, the first run starts a detached server (``connect/local_server.py``) and + # records it in a discovery file, and later runs reconnect to it instead of re-paying the cold + # start. + + @staticmethod + def _local_connect_discovery_path() -> str: + """Location of the discovery file describing the running persistent local server.""" + override = os.environ.get("SPARK_LOCAL_CONNECT_DISCOVERY") + if override: + return override + return os.path.join(os.path.expanduser("~"), ".spark", "connect-local.json") + + @staticmethod + def _read_local_connect_discovery() -> Optional[Dict[str, Any]]: + """Read and validate the discovery file, returning ``None`` if it is absent or malformed.""" + path = SparkSession._local_connect_discovery_path() + try: + with open(path, "r") as f: + disc = json.load(f) + except (OSError, ValueError): + return None + if not isinstance(disc, dict) or not all( + k in disc for k in ("host", "port", "token", "pid", "spark_version") + ): + return None + return disc + + @staticmethod + def _local_connect_server_is_reusable(disc: Dict[str, Any]) -> bool: + """Decide whether the server described by ``disc`` can be reused by this process. + + Reuse requires that the recorded Spark version matches this client's, the recorded process + is still alive, and it is accepting connections on the recorded port. A version mismatch, + dead pid, or closed port means we must start our own server instead. + """ + import socket + from pyspark.version import __version__ + + if disc.get("spark_version") != __version__: Review Comment: It's weird that we use `get` here and just `[]` later. I think we already validate the dict when we read it so we don't need to do `get`. ########## python/pyspark/sql/connect/session.py: ########## @@ -1247,6 +1249,372 @@ def _start_connect_server(master: str, opts: Dict[str, Any]) -> None: messageParameters={}, ) + # Opt-in reuse of a persistent local Spark Connect server. By default ``.remote("local[*]")`` + # boots a fresh in-process server every process (see ``_start_connect_server`` above); when + # reuse is enabled, the first run starts a detached server (``connect/local_server.py``) and + # records it in a discovery file, and later runs reconnect to it instead of re-paying the cold + # start. + + @staticmethod + def _local_connect_discovery_path() -> str: + """Location of the discovery file describing the running persistent local server.""" + override = os.environ.get("SPARK_LOCAL_CONNECT_DISCOVERY") + if override: + return override + return os.path.join(os.path.expanduser("~"), ".spark", "connect-local.json") + + @staticmethod + def _read_local_connect_discovery() -> Optional[Dict[str, Any]]: + """Read and validate the discovery file, returning ``None`` if it is absent or malformed.""" + path = SparkSession._local_connect_discovery_path() + try: + with open(path, "r") as f: + disc = json.load(f) + except (OSError, ValueError): + return None + if not isinstance(disc, dict) or not all( + k in disc for k in ("host", "port", "token", "pid", "spark_version") + ): + return None + return disc + + @staticmethod + def _local_connect_server_is_reusable(disc: Dict[str, Any]) -> bool: + """Decide whether the server described by ``disc`` can be reused by this process. + + Reuse requires that the recorded Spark version matches this client's, the recorded process + is still alive, and it is accepting connections on the recorded port. A version mismatch, + dead pid, or closed port means we must start our own server instead. + """ + import socket + from pyspark.version import __version__ + + if disc.get("spark_version") != __version__: + return False + try: + os.kill(int(disc["pid"]), 0) + except (ProcessLookupError, ValueError, TypeError): + return False + except OSError: + # The process exists but is not ours to signal (e.g. PermissionError) -- treat as alive. + pass + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(0.5) + if sock.connect_ex((disc["host"], int(disc["port"]))) != 0: + return False + return True + + @staticmethod + def _reuse_from_discovery() -> Optional[str]: + """Return an endpoint for the recorded server if it is reusable, else ``None``. + + On success it also sets ``SPARK_CONNECT_AUTHENTICATE_TOKEN`` so the client authenticates + against that server. + """ + disc = SparkSession._read_local_connect_discovery() + if disc is not None and SparkSession._local_connect_server_is_reusable(disc): + os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = disc["token"] + return "sc://{}:{}".format(disc["host"], disc["port"]) + return None + + @staticmethod + def _acquire_local_connect_start_lock() -> Any: + """Take an exclusive file lock guarding persistent-server start-up. + + Returns the open fd to pass to ``_release_local_connect_start_lock``. Returns ``None`` on + platforms without ``fcntl``; those callers may race and then reconnect to the server that + wins the discovery-file update. + """ + try: + import fcntl + except ImportError: + return None + path = SparkSession._local_connect_discovery_path() + ".lock" + parent = os.path.dirname(path) + if parent and not os.path.isdir(parent): + os.makedirs(parent, exist_ok=True) + fd = os.open(path, os.O_RDWR | os.O_CREAT, 0o600) + fcntl.flock(fd, fcntl.LOCK_EX) + return fd + + @staticmethod + def _release_local_connect_start_lock(fd: Any) -> None: + if fd is None: + return + try: + import fcntl + + fcntl.flock(fd, fcntl.LOCK_UN) + finally: + os.close(fd) + + @staticmethod + def _reuse_or_start_local_connect_server(master: str, opts: Dict[str, Any]) -> str: + """Reuse a running persistent local Connect server, or start one if none is reusable. + + Returns the ``sc://host:port`` endpoint to connect to. This is the opt-in counterpart of + ``_start_connect_server`` and is only reached for a ``local`` master when + ``spark.local.connect.reuse`` / ``SPARK_LOCAL_CONNECT_REUSE`` is set. + """ + # Fast path: reuse an already-running server without taking the cross-process lock. + endpoint = SparkSession._reuse_from_discovery() + if endpoint is not None: + return endpoint + # No reusable server yet. Serialize start-up across processes when file locking is + # available. Without it, racing callers may each start a daemon; the winner writes the + # discovery file and the others reconnect to it. + lock_fd = SparkSession._acquire_local_connect_start_lock() + try: + endpoint = SparkSession._reuse_from_discovery() + if endpoint is not None: + return endpoint + return SparkSession._start_persistent_local_connect_server(master, opts) + finally: + SparkSession._release_local_connect_start_lock(lock_fd) + + @staticmethod + def _local_port_available(port: int) -> bool: + """Whether ``port`` can currently be bound on localhost (best effort, subject to races).""" + import socket + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + try: + sock.bind(("localhost", port)) + return True + except OSError: + return False + + @staticmethod + def _local_connect_server_conf(opts: Dict[str, Any]) -> Dict[str, Any]: + """Start-up configs to seed a freshly started persistent server. + + Mirrors the merge that the in-process ``_start_connect_server`` applies to its ``SparkConf`` + so that first-run behavior matches (warehouse dir, app name, jars/packages, catalog confs, + etc.). Keys the daemon controls itself (master, port, token, plugins) and the reuse opt-in + keys are excluded. This only seeds the run that *starts* the server; a later run + reconnecting to an already-warm JVM cannot change its static configs. + """ + conf: Dict[str, Any] = {} + for i in range(int(os.environ.get("PYSPARK_REMOTE_INIT_CONF_LEN", "0"))): + conf = json.loads(os.environ["PYSPARK_REMOTE_INIT_CONF_{}".format(i)]) + conf.update(opts) + for k in ( + "spark.remote", + "spark.api.mode", + "spark.master", + "spark.connect.authenticate.token", + "spark.connect.grpc.binding.port", + "spark.local.connect.reuse", + "spark.local.connect.server.port", + "spark.local.connect.server.idleTimeout", + ): + conf.pop(k, None) + return conf + + @staticmethod + def _signal_local_connect_server(pid: int, sig: int) -> bool: + """Signal a detached local Connect daemon, including its JVM on POSIX. + + The daemon is launched as a session leader, so its process group id equals its pid and + signalling the group reaps its child JVM too. If the group id differs, ``pid`` belongs to + an unrelated process (e.g. recycled after a stale discovery file), so only the pid itself + is signalled -- never a group this code did not create. + """ + try: + if os.name == "posix" and os.getpgid(pid) == pid: + os.killpg(pid, sig) + else: + os.kill(pid, sig) + return True + except OSError: + return False + + @staticmethod + def _terminate_local_connect_server(proc: Any) -> None: + """Terminate a daemon started by ``_start_persistent_local_connect_server`` and its JVM. + + The group SIGTERM asks the daemon and its JVM to shut down gracefully, but the daemon + exiting does not mean the JVM is gone: its graceful shutdown can outlive the daemon. On + POSIX this therefore waits for the whole process group to disappear and escalates to a + group SIGKILL if any member outlives the grace period. + """ + if not SparkSession._signal_local_connect_server(proc.pid, signal.SIGTERM): + return + try: + proc.wait(timeout=10) + except Exception: + pass + if os.name != "posix": + if proc.poll() is None: + proc.kill() + return + # The group id stays valid while any member (i.e. the JVM) is alive, even after the + # daemon leader has been reaped, so poll and signal the group id directly rather than + # via _signal_local_connect_server (whose getpgid lookup needs a live leader). Signalling + # this group is safe: it was created above us by this launch, not read from disk. + deadline = time.time() + 10 + while time.time() < deadline: + proc.poll() # reap the daemon if it exited, so only live members keep the group + try: + os.killpg(proc.pid, 0) + except OSError: + return + time.sleep(0.2) + try: + os.killpg(proc.pid, signal.SIGKILL) + except OSError: + pass + proc.poll() Review Comment: What does this do? ########## python/pyspark/sql/connect/session.py: ########## @@ -1247,6 +1249,372 @@ def _start_connect_server(master: str, opts: Dict[str, Any]) -> None: messageParameters={}, ) + # Opt-in reuse of a persistent local Spark Connect server. By default ``.remote("local[*]")`` + # boots a fresh in-process server every process (see ``_start_connect_server`` above); when + # reuse is enabled, the first run starts a detached server (``connect/local_server.py``) and + # records it in a discovery file, and later runs reconnect to it instead of re-paying the cold + # start. + + @staticmethod + def _local_connect_discovery_path() -> str: + """Location of the discovery file describing the running persistent local server.""" + override = os.environ.get("SPARK_LOCAL_CONNECT_DISCOVERY") Review Comment: Isn't this just ```python os.environ.get( "SPARK_LOCAL_CONNECT_DISCOVERY", os.path.join(os.path.expanduser("~"), ".spark", "connect-local.json") ) ``` Or you are protecting something like empty string? ########## python/pyspark/sql/connect/local_server.py: ########## @@ -0,0 +1,212 @@ +# +# 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. +# + +""" +Long-lived local Spark Connect server for the opt-in reuse path. + +This module is launched as a *detached* child process by +``SparkSession._reuse_or_start_local_connect_server`` (in ``pyspark.sql.connect.session``) when +``SPARK_LOCAL_CONNECT_REUSE`` / ``spark.local.connect.reuse`` is enabled. It starts a regular +(classic) local Spark session with the Spark Connect plugin -- the same mechanism the in-process +``SparkSession._start_connect_server`` uses -- and then blocks, so one warm JVM and Spark Connect +server keeps serving many short-lived client processes. Each client connection gets its own +isolated server-side session, so session-local state (temp views, runtime SQL confs, isolated +artifacts) does not leak between runs. + +Once the server is accepting connections it writes a discovery file (host, the actually bound port, +the auth token, its pid and the Spark version) that later client processes read to reconnect. + +It is launched by file path rather than ``python -m`` so it does not require the Spark Connect +*client* dependencies (grpc, etc.): a server only needs a classic PySpark install plus the Connect +server jar, like ``sbin/start-connect-server.sh``. It imports only the classic ``pyspark.sql`` API. +""" + +import sys + +# Launching by file path puts this file's directory -- pyspark/sql/connect -- at the front of +# sys.path, where modules such as `types` and `logging` shadow the standard library and break +# ordinary imports. Drop it before importing anything else (`import sys` cannot be shadowed); +# pyspark stays importable via the remaining sys.path entries. +if sys.path: + del sys.path[0] Review Comment: Things like this is why I'm against running it as a file. ########## python/pyspark/sql/connect/local_server.py: ########## @@ -0,0 +1,212 @@ +# +# 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. +# + +""" +Long-lived local Spark Connect server for the opt-in reuse path. + +This module is launched as a *detached* child process by +``SparkSession._reuse_or_start_local_connect_server`` (in ``pyspark.sql.connect.session``) when +``SPARK_LOCAL_CONNECT_REUSE`` / ``spark.local.connect.reuse`` is enabled. It starts a regular +(classic) local Spark session with the Spark Connect plugin -- the same mechanism the in-process +``SparkSession._start_connect_server`` uses -- and then blocks, so one warm JVM and Spark Connect +server keeps serving many short-lived client processes. Each client connection gets its own +isolated server-side session, so session-local state (temp views, runtime SQL confs, isolated +artifacts) does not leak between runs. + +Once the server is accepting connections it writes a discovery file (host, the actually bound port, +the auth token, its pid and the Spark version) that later client processes read to reconnect. + +It is launched by file path rather than ``python -m`` so it does not require the Spark Connect +*client* dependencies (grpc, etc.): a server only needs a classic PySpark install plus the Connect +server jar, like ``sbin/start-connect-server.sh``. It imports only the classic ``pyspark.sql`` API. +""" + +import sys + +# Launching by file path puts this file's directory -- pyspark/sql/connect -- at the front of +# sys.path, where modules such as `types` and `logging` shadow the standard library and break +# ordinary imports. Drop it before importing anything else (`import sys` cannot be shadowed); +# pyspark stays importable via the remaining sys.path entries. +if sys.path: + del sys.path[0] + +import argparse +import json +import os +import signal +import time +from typing import Any + + +def _write_discovery(path: str, host: str, port: int, token: str, version: str) -> None: + """Atomically write the discovery file with ``0600`` perms (it holds the auth token).""" + parent = os.path.dirname(path) + if parent and not os.path.isdir(parent): Review Comment: What's this line protect against? ########## python/pyspark/sql/connect/local_server.py: ########## @@ -0,0 +1,212 @@ +# +# 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. +# + +""" +Long-lived local Spark Connect server for the opt-in reuse path. + +This module is launched as a *detached* child process by +``SparkSession._reuse_or_start_local_connect_server`` (in ``pyspark.sql.connect.session``) when +``SPARK_LOCAL_CONNECT_REUSE`` / ``spark.local.connect.reuse`` is enabled. It starts a regular +(classic) local Spark session with the Spark Connect plugin -- the same mechanism the in-process +``SparkSession._start_connect_server`` uses -- and then blocks, so one warm JVM and Spark Connect +server keeps serving many short-lived client processes. Each client connection gets its own +isolated server-side session, so session-local state (temp views, runtime SQL confs, isolated +artifacts) does not leak between runs. + +Once the server is accepting connections it writes a discovery file (host, the actually bound port, +the auth token, its pid and the Spark version) that later client processes read to reconnect. + +It is launched by file path rather than ``python -m`` so it does not require the Spark Connect +*client* dependencies (grpc, etc.): a server only needs a classic PySpark install plus the Connect +server jar, like ``sbin/start-connect-server.sh``. It imports only the classic ``pyspark.sql`` API. +""" + +import sys + +# Launching by file path puts this file's directory -- pyspark/sql/connect -- at the front of +# sys.path, where modules such as `types` and `logging` shadow the standard library and break +# ordinary imports. Drop it before importing anything else (`import sys` cannot be shadowed); +# pyspark stays importable via the remaining sys.path entries. +if sys.path: + del sys.path[0] + +import argparse +import json +import os +import signal +import time +from typing import Any + + +def _write_discovery(path: str, host: str, port: int, token: str, version: str) -> None: + """Atomically write the discovery file with ``0600`` perms (it holds the auth token).""" + parent = os.path.dirname(path) + if parent and not os.path.isdir(parent): + os.makedirs(parent, exist_ok=True) + payload = { + "host": host, + "port": port, + "token": token, + "pid": os.getpid(), + "spark_version": version, + } + tmp = "{}.{}.tmp".format(path, os.getpid()) + fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as f: + f.write(json.dumps(payload)) + os.replace(tmp, path) + + +def _remove_discovery_if_ours(path: str) -> None: + """Remove the discovery file, but only if it still points at this process.""" + try: + with open(path, "r") as f: + disc = json.load(f) + except (OSError, ValueError): + return + if disc.get("pid") == os.getpid(): Review Comment: Do we care about any racing issues? ########## python/pyspark/sql/connect/local_server.py: ########## @@ -0,0 +1,212 @@ +# +# 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. +# + +""" +Long-lived local Spark Connect server for the opt-in reuse path. + +This module is launched as a *detached* child process by +``SparkSession._reuse_or_start_local_connect_server`` (in ``pyspark.sql.connect.session``) when +``SPARK_LOCAL_CONNECT_REUSE`` / ``spark.local.connect.reuse`` is enabled. It starts a regular +(classic) local Spark session with the Spark Connect plugin -- the same mechanism the in-process +``SparkSession._start_connect_server`` uses -- and then blocks, so one warm JVM and Spark Connect +server keeps serving many short-lived client processes. Each client connection gets its own +isolated server-side session, so session-local state (temp views, runtime SQL confs, isolated +artifacts) does not leak between runs. + +Once the server is accepting connections it writes a discovery file (host, the actually bound port, +the auth token, its pid and the Spark version) that later client processes read to reconnect. + +It is launched by file path rather than ``python -m`` so it does not require the Spark Connect +*client* dependencies (grpc, etc.): a server only needs a classic PySpark install plus the Connect +server jar, like ``sbin/start-connect-server.sh``. It imports only the classic ``pyspark.sql`` API. +""" + +import sys + +# Launching by file path puts this file's directory -- pyspark/sql/connect -- at the front of +# sys.path, where modules such as `types` and `logging` shadow the standard library and break +# ordinary imports. Drop it before importing anything else (`import sys` cannot be shadowed); +# pyspark stays importable via the remaining sys.path entries. +if sys.path: + del sys.path[0] + +import argparse +import json +import os +import signal +import time +from typing import Any + + +def _write_discovery(path: str, host: str, port: int, token: str, version: str) -> None: + """Atomically write the discovery file with ``0600`` perms (it holds the auth token).""" + parent = os.path.dirname(path) + if parent and not os.path.isdir(parent): + os.makedirs(parent, exist_ok=True) + payload = { + "host": host, + "port": port, + "token": token, + "pid": os.getpid(), + "spark_version": version, + } + tmp = "{}.{}.tmp".format(path, os.getpid()) + fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as f: + f.write(json.dumps(payload)) + os.replace(tmp, path) + + +def _remove_discovery_if_ours(path: str) -> None: + """Remove the discovery file, but only if it still points at this process.""" + try: + with open(path, "r") as f: + disc = json.load(f) + except (OSError, ValueError): + return + if disc.get("pid") == os.getpid(): + try: + os.remove(path) + except OSError: + pass + + +def _has_active_sessions(spark: Any) -> bool: + """Best-effort check for whether any Spark Connect session is currently registered. + + Used only by the idle-shutdown reaper. Any failure (server not started yet, API drift, py4j + error) returns ``True`` so the reaper never terminates a server it cannot inspect. + """ + jvm = spark.sparkContext._jvm + service = getattr( + getattr(jvm, "org.apache.spark.sql.connect.service.SparkConnectService$"), + "MODULE$", + ) + if not service.started(): + return True + return not service.sessionManager().listActiveSessions().isEmpty() + + +def _bound_port(spark: Any, requested_port: int) -> int: + """Return the port the Connect server actually bound (``requested_port`` may have been 0).""" + jvm = spark.sparkContext._jvm + service = getattr( + getattr(jvm, "org.apache.spark.sql.connect.service.SparkConnectService$"), + "MODULE$", + ) + try: + return int(service.localPort()) + except Exception: + return requested_port + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--master", default="local[*]") + parser.add_argument("--port", type=int, default=15002) + parser.add_argument("--token", default=None) + parser.add_argument("--discovery", required=True) + parser.add_argument( + "--conf-file", + default=None, + help="path to a JSON file of extra SparkConf entries to seed the server with", + ) + parser.add_argument( + "--idle-timeout", + type=float, + default=3600.0, + help="seconds with no active session after which the server self-terminates; <=0 disables", + ) + parser.add_argument("--poll-interval", type=float, default=60.0) + args = parser.parse_args() + + # Build a CLASSIC session: the connect-mode env vars would otherwise divert us into a client. + for var in ("SPARK_REMOTE", "SPARK_LOCAL_REMOTE", "SPARK_CONNECT_MODE_ENABLED"): + os.environ.pop(var, None) + if args.token: + os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = args.token + + from pyspark.sql import SparkSession + from pyspark.version import __version__ + + builder = SparkSession.builder.master(args.master) + # Seed the caller's start-up confs first so first-run behavior matches the in-process path, then + # apply the settings the server itself controls so they always win. + if args.conf_file: + with open(args.conf_file, "r") as f: + for key, value in json.load(f).items(): + builder = builder.config(key, str(value)) + builder = ( + builder.config("spark.plugins", "org.apache.spark.sql.connect.SparkConnectPlugin") + .config("spark.connect.grpc.binding.port", str(args.port)) + # Match the isolation the in-process `.remote("local[*]")` path sets so per-session + # artifacts (added jars/files/classes) do not leak across the sessions this server hosts. + .config("spark.sql.artifact.isolation.enabled", "true") + .config("spark.sql.artifact.isolation.alwaysApplyClassloader", "true") + ) + if args.token: + builder = builder.config("spark.connect.authenticate.token", args.token) + spark = builder.getOrCreate() + + bound_port = _bound_port(spark, args.port) + _write_discovery(args.discovery, "localhost", bound_port, args.token, __version__) + print( + "SPARK-CONNECT-LOCAL-SERVER READY port={} pid={}".format(bound_port, os.getpid()), + flush=True, + ) + + stop = {"flag": False} + + def _handle(_signum: int, _frame: Any) -> None: + stop["flag"] = True + + signal.signal(signal.SIGTERM, _handle) + try: + signal.signal(signal.SIGINT, _handle) + except ValueError: + # SIGINT may not be settable when not on the main thread on some platforms. Review Comment: Why would this be set on threads that are not main? It's in `main()`. It's not about `SIGINT`, `signal.signal()` should only be used in main thread. ########## python/pyspark/sql/connect/session.py: ########## @@ -1247,6 +1249,372 @@ def _start_connect_server(master: str, opts: Dict[str, Any]) -> None: messageParameters={}, ) + # Opt-in reuse of a persistent local Spark Connect server. By default ``.remote("local[*]")`` + # boots a fresh in-process server every process (see ``_start_connect_server`` above); when + # reuse is enabled, the first run starts a detached server (``connect/local_server.py``) and + # records it in a discovery file, and later runs reconnect to it instead of re-paying the cold + # start. + + @staticmethod Review Comment: If every new method added is a `@staticmethod`, it means the code does not belong to this class. ########## python/pyspark/sql/connect/session.py: ########## @@ -1247,6 +1249,372 @@ def _start_connect_server(master: str, opts: Dict[str, Any]) -> None: messageParameters={}, ) + # Opt-in reuse of a persistent local Spark Connect server. By default ``.remote("local[*]")`` + # boots a fresh in-process server every process (see ``_start_connect_server`` above); when + # reuse is enabled, the first run starts a detached server (``connect/local_server.py``) and + # records it in a discovery file, and later runs reconnect to it instead of re-paying the cold + # start. + + @staticmethod + def _local_connect_discovery_path() -> str: + """Location of the discovery file describing the running persistent local server.""" + override = os.environ.get("SPARK_LOCAL_CONNECT_DISCOVERY") + if override: + return override + return os.path.join(os.path.expanduser("~"), ".spark", "connect-local.json") + + @staticmethod + def _read_local_connect_discovery() -> Optional[Dict[str, Any]]: + """Read and validate the discovery file, returning ``None`` if it is absent or malformed.""" + path = SparkSession._local_connect_discovery_path() + try: + with open(path, "r") as f: + disc = json.load(f) + except (OSError, ValueError): + return None + if not isinstance(disc, dict) or not all( + k in disc for k in ("host", "port", "token", "pid", "spark_version") + ): + return None + return disc + + @staticmethod + def _local_connect_server_is_reusable(disc: Dict[str, Any]) -> bool: + """Decide whether the server described by ``disc`` can be reused by this process. + + Reuse requires that the recorded Spark version matches this client's, the recorded process + is still alive, and it is accepting connections on the recorded port. A version mismatch, + dead pid, or closed port means we must start our own server instead. + """ + import socket + from pyspark.version import __version__ + + if disc.get("spark_version") != __version__: + return False + try: + os.kill(int(disc["pid"]), 0) Review Comment: I don't think this is a no-op on Windows. ########## python/pyspark/sql/connect/session.py: ########## @@ -1247,6 +1249,372 @@ def _start_connect_server(master: str, opts: Dict[str, Any]) -> None: messageParameters={}, ) + # Opt-in reuse of a persistent local Spark Connect server. By default ``.remote("local[*]")`` + # boots a fresh in-process server every process (see ``_start_connect_server`` above); when + # reuse is enabled, the first run starts a detached server (``connect/local_server.py``) and + # records it in a discovery file, and later runs reconnect to it instead of re-paying the cold + # start. + + @staticmethod + def _local_connect_discovery_path() -> str: + """Location of the discovery file describing the running persistent local server.""" + override = os.environ.get("SPARK_LOCAL_CONNECT_DISCOVERY") + if override: + return override + return os.path.join(os.path.expanduser("~"), ".spark", "connect-local.json") + + @staticmethod + def _read_local_connect_discovery() -> Optional[Dict[str, Any]]: + """Read and validate the discovery file, returning ``None`` if it is absent or malformed.""" + path = SparkSession._local_connect_discovery_path() + try: + with open(path, "r") as f: + disc = json.load(f) + except (OSError, ValueError): + return None + if not isinstance(disc, dict) or not all( + k in disc for k in ("host", "port", "token", "pid", "spark_version") + ): + return None + return disc + + @staticmethod + def _local_connect_server_is_reusable(disc: Dict[str, Any]) -> bool: + """Decide whether the server described by ``disc`` can be reused by this process. + + Reuse requires that the recorded Spark version matches this client's, the recorded process + is still alive, and it is accepting connections on the recorded port. A version mismatch, + dead pid, or closed port means we must start our own server instead. + """ + import socket + from pyspark.version import __version__ + + if disc.get("spark_version") != __version__: + return False + try: + os.kill(int(disc["pid"]), 0) + except (ProcessLookupError, ValueError, TypeError): + return False + except OSError: + # The process exists but is not ours to signal (e.g. PermissionError) -- treat as alive. + pass + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(0.5) + if sock.connect_ex((disc["host"], int(disc["port"]))) != 0: + return False + return True + + @staticmethod + def _reuse_from_discovery() -> Optional[str]: + """Return an endpoint for the recorded server if it is reusable, else ``None``. + + On success it also sets ``SPARK_CONNECT_AUTHENTICATE_TOKEN`` so the client authenticates + against that server. + """ + disc = SparkSession._read_local_connect_discovery() + if disc is not None and SparkSession._local_connect_server_is_reusable(disc): + os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = disc["token"] + return "sc://{}:{}".format(disc["host"], disc["port"]) + return None + + @staticmethod + def _acquire_local_connect_start_lock() -> Any: Review Comment: Let's change this lock to a context pattern with `@contextlib.contextmanager` and use it with `with`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
