From 7374dc2f1aae3fff06706cd77bb2f704a279448f Mon Sep 17 00:00:00 2001 From: ericm-db Date: Tue, 30 Jun 2026 17:46:24 +0000 Subject: [PATCH 01/14] [SPARK-57787][CONNECT] Reuse a persistent local Spark Connect server for faster local startup Adds an opt-in (SPARK_LOCAL_CONNECT_REUSE / spark.local.connect.reuse) so that `SparkSession.builder.remote("local[*]").getOrCreate()` reconnects to a persistent local Spark Connect server -- starting a detached one on the first run and reconnecting to it on later runs -- instead of booting a fresh in-process server in every process. The first run pays the cold start once; later runs reconnect in a fraction of a second. Default behavior is unchanged when the opt-in is off. No protocol or Scala changes. --- dev/sparktestsupport/modules.py | 1 + docs/spark-connect-overview.md | 62 ++++++ python/pyspark/errors/error-conditions.json | 5 + python/pyspark/sql/connect/local_server.py | 205 ++++++++++++++++++ python/pyspark/sql/connect/session.py | 194 +++++++++++++++++ python/pyspark/sql/session.py | 22 +- .../connect/test_connect_local_server.py | 199 +++++++++++++++++ 7 files changed, 687 insertions(+), 1 deletion(-) create mode 100644 python/pyspark/sql/connect/local_server.py create mode 100644 python/pyspark/sql/tests/connect/test_connect_local_server.py diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index 860158b941f64..e528b7fda0399 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -1162,6 +1162,7 @@ def __hash__(self): "pyspark.sql.tests.connect.test_connect_readwriter", "pyspark.sql.tests.connect.test_connect_retry", "pyspark.sql.tests.connect.test_connect_session", + "pyspark.sql.tests.connect.test_connect_local_server", "pyspark.sql.tests.connect.test_connect_stat", "pyspark.sql.tests.connect.test_parity_geographytype", "pyspark.sql.tests.connect.test_parity_geometrytype", diff --git a/docs/spark-connect-overview.md b/docs/spark-connect-overview.md index 3c15153e03053..f0ee4f54db2de 100644 --- a/docs/spark-connect-overview.md +++ b/docs/spark-connect-overview.md @@ -277,6 +277,68 @@ The connection may also be programmatically created using _SparkSession#builder_ +## Faster local iteration with a persistent Connect server + +When you develop or test locally with + +{% highlight python %} +from pyspark.sql import SparkSession +spark = SparkSession.builder.remote("local[*]").getOrCreate() +{% endhighlight %} + +PySpark boots a fresh in-process Spark Connect server in **every** process. Each +`python script.py` run (or each forked test JVM) therefore re-pays the one-time startup cost -- +JVM warmup, `SparkContext` construction, and Connect server boot -- which can take a few seconds and +makes a quick edit/run loop feel slow. + +There are two ways to amortize that cost across runs by reconnecting to a long-lived local server. + +### Start a server yourself and connect to it + +Start one persistent local Spark Connect server and point every run at it: + +{% highlight bash %} +# Start once; it stays up across runs. +$SPARK_HOME/sbin/start-connect-server.sh --master "local[*]" + +# Every run reconnects instead of booting a new server. +python -c 'from pyspark.sql import SparkSession; SparkSession.builder.remote("sc://localhost:15002").getOrCreate()' + +# Stop it when you are done. +$SPARK_HOME/sbin/stop-connect-server.sh +{% endhighlight %} + +### Let PySpark manage the server (opt-in) + +If you would rather keep your code as `SparkSession.builder.remote("local[*]").getOrCreate()` and not +manage a server by hand, enable the opt-in reuse path. The first run starts a **detached** local +Connect server and records it in a discovery file; later runs reconnect to it in a fraction of a +second: + +{% highlight bash %} +export SPARK_LOCAL_CONNECT_REUSE=1 # or .config("spark.local.connect.reuse", "true") +python script.py # 1st run: starts a persistent server (cold start, once) +python script.py # 2nd+ run: reconnects to it (sub-second) +{% endhighlight %} + +This is **off by default**; nothing changes unless you opt in. A few details: + +- Each run connects as its own Connect session, so session-local state -- temp views, runtime SQL + configurations, and (with artifact isolation, which stays on) session artifacts -- is fresh on + every run and never leaks between runs. State backed by the shared `SparkContext` (the persistent + catalog/warehouse, global temp views, and cached datasets) *is* shared across runs, so namespace + per-run databases or clear that state yourself if your runs must be fully isolated. +- The server listens on port `15002` by default and authenticates with a token written, together + with the host, port, pid and Spark version, to `~/.spark/connect-local.json` (mode `0600`). + Set `SPARK_LOCAL_CONNECT_DISCOVERY` to relocate that file. Reuse is refused (and a fresh server + started) if the recorded process is gone, the port is closed, or the Spark version differs. +- The server self-terminates after it has been idle for `spark.local.connect.server.idleTimeout` + seconds (default `3600`; set `0` to disable). To stop it explicitly: + +{% highlight bash %} +python -c "from pyspark.sql.connect.session import SparkSession; SparkSession._stop_local_connect_server()" +{% endhighlight %} + ## Use Spark Connect in standalone applications
diff --git a/python/pyspark/errors/error-conditions.json b/python/pyspark/errors/error-conditions.json index 38417cbf01889..02eb4aa685d40 100644 --- a/python/pyspark/errors/error-conditions.json +++ b/python/pyspark/errors/error-conditions.json @@ -546,6 +546,11 @@ " and should be of the same length, got and ." ] }, + "LOCAL_CONNECT_SERVER_START_FAILED": { + "message": [ + "Failed to start a persistent local Spark Connect server: ." + ] + }, "LOCAL_RELATION_SIZE_LIMIT_EXCEEDED": { "message": [ "Local relation size ( bytes) exceeds the limit ( bytes)." diff --git a/python/pyspark/sql/connect/local_server.py b/python/pyspark/sql/connect/local_server.py new file mode 100644 index 0000000000000..a5d5298eb3662 --- /dev/null +++ b/python/pyspark/sql/connect/local_server.py @@ -0,0 +1,205 @@ +# +# 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 (not ``python -m pyspark.sql.connect.local_server``) on purpose: it only +needs a classic PySpark install plus the Spark Connect server jar on the classpath -- exactly like +``sbin/start-connect-server.sh`` -- and must not require the Spark Connect *client* dependencies +(grpc, etc.). Importing the ``pyspark.sql.connect`` package would trigger that client-dependency +check, so this file imports only the classic ``pyspark.sql`` API. +""" +import sys + +# This module is launched by file path (see SparkSession._start_persistent_local_connect_server), +# which puts its own directory -- pyspark/sql/connect -- at the front of sys.path. That directory +# holds modules whose names collide with the standard library (e.g. `types`, `logging`), so leaving +# it there shadows the stdlib and breaks ordinary imports. Drop it before importing anything else; +# pyspark itself stays importable via the remaining sys.path entries (PYTHONPATH / installation). +# `import sys` is a built-in and cannot be shadowed, so it is safe to run first. +if sys.path: + del sys.path[0] + +import argparse # noqa: E402 +import json # noqa: E402 +import os # noqa: E402 +import signal # noqa: E402 +import time # noqa: E402 +from typing import Any # noqa: E402 + + +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( + "--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) + .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__) + # Printed to the (normally discarded) child stdout; useful when launched with stdout attached. + 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. + pass + + idle_timeout = args.idle_timeout + poll_interval = max(1.0, args.poll_interval) + last_active = time.monotonic() + last_poll = 0.0 + try: + # Sleep in short slices so a SIGTERM is observed promptly regardless of the poll interval. + while not stop["flag"]: + time.sleep(1.0) + if idle_timeout <= 0: + continue + now = time.monotonic() + if now - last_poll < poll_interval: + continue + last_poll = now + try: + active = _has_active_sessions(spark) + except Exception: + active = True # fail open: never reap a server we cannot inspect + if active: + last_active = now + elif now - last_active > idle_timeout: + break + finally: + _remove_discovery_if_ours(args.discovery) + spark.stop() + + +if __name__ == "__main__": + main() diff --git a/python/pyspark/sql/connect/session.py b/python/pyspark/sql/connect/session.py index 99da7308a1a00..2542776d0df0c 100644 --- a/python/pyspark/sql/connect/session.py +++ b/python/pyspark/sql/connect/session.py @@ -17,9 +17,11 @@ import uuid import json +import signal import threading import os import sys +import time import warnings from collections.abc import Callable, Sized import functools @@ -1247,6 +1249,198 @@ def _start_connect_server(master: str, opts: Dict[str, Any]) -> None: messageParameters={}, ) + # --------------------------------------------------------------------------------------------- + # Opt-in reuse of a persistent local Spark Connect server. + # + # By default ``SparkSession.builder.remote("local[*]")`` boots a fresh in-process Connect server + # in every process (see ``_start_connect_server`` above), so every ``python script.py`` run + # re-pays the cold start (JVM warmup + SparkContext + server boot). When the opt-in is enabled, + # the first run starts a *detached* server (``connect/local_server.py``) and records it in a + # discovery file; later runs reconnect to it in a fraction of a second instead. + # --------------------------------------------------------------------------------------------- + + @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: + return False + except (TypeError, ValueError): + return False + except OSError: + # e.g. PermissionError: the process exists but is not ours to signal -- 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_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 and sets + ``SPARK_CONNECT_AUTHENTICATE_TOKEN`` so the client authenticates against the server. 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. + """ + 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 SparkSession._start_persistent_local_connect_server(master, opts) + + @staticmethod + def _start_persistent_local_connect_server(master: str, opts: Dict[str, Any]) -> str: + """Launch a detached persistent local Connect server and wait until it is reachable.""" + import socket + import subprocess + + discovery_path = SparkSession._local_connect_discovery_path() + token = opts.get("spark.connect.authenticate.token") or str(uuid.uuid4()) + + # Tests use an ephemeral port (0) so they can run in parallel; the server reports the actual + # bound port back through the discovery file. + if "SPARK_TESTING" in os.environ: + port = 0 + else: + port = int( + opts.get("spark.local.connect.server.port", DefaultChannelBuilder.default_port()) + ) + idle_timeout = opts.get("spark.local.connect.server.idleTimeout", "3600") + + daemon = os.path.join(os.path.dirname(os.path.abspath(__file__)), "local_server.py") + cmd = [ + sys.executable, + daemon, + "--master", + master, + "--port", + str(port), + "--token", + token, + "--discovery", + discovery_path, + "--idle-timeout", + str(idle_timeout), + ] + + # Launch detached so the server outlives this client process. + env = dict(os.environ) + for var in ("SPARK_REMOTE", "SPARK_LOCAL_REMOTE", "SPARK_CONNECT_MODE_ENABLED"): + env.pop(var, None) + popen_kwargs: Dict[str, Any] = dict( + env=env, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + if os.name == "posix": + popen_kwargs["start_new_session"] = True + else: + detached = getattr(subprocess, "DETACHED_PROCESS", 0) + new_group = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0) + popen_kwargs["creationflags"] = detached | new_group + proc = subprocess.Popen(cmd, **popen_kwargs) + + # Wait for the server to write the discovery file (which records the actual bound port) and + # start accepting connections. + deadline = time.time() + 120 + while time.time() < deadline: + exit_code = proc.poll() + if exit_code is not None: + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={ + "reason": "the server process exited with code {}".format(exit_code) + }, + ) + disc = SparkSession._read_local_connect_discovery() + if disc is not None and disc.get("pid") == proc.pid and disc.get("token") == token: + 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: + os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = token + return "sc://{}:{}".format(disc["host"], disc["port"]) + time.sleep(0.25) + + # Timed out: do not leave the detached process orphaned. + try: + proc.terminate() + except OSError: + pass + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={"reason": "the server did not become ready within 120 seconds"}, + ) + + @staticmethod + def _stop_local_connect_server() -> bool: + """Stop the persistent local Connect server started by the reuse path, if any. + + Returns ``True`` if a running server was signalled to stop. Safe to call when none is + running. Intended for the dev/test loop, e.g.:: + + python -c "from pyspark.sql.connect.session import SparkSession; \\ + SparkSession._stop_local_connect_server()" + """ + disc = SparkSession._read_local_connect_discovery() + path = SparkSession._local_connect_discovery_path() + stopped = False + if disc is not None: + try: + pid = int(disc["pid"]) + except (TypeError, ValueError, KeyError): + pid = None + # Never signal the current process: a discovery file should only ever point at the + # detached server, never at the client that is reading it. + if pid is not None and pid != os.getpid(): + try: + os.kill(pid, signal.SIGTERM) + stopped = True + except OSError: + pass + try: + os.remove(path) + except OSError: + pass + return stopped + @property def session_id(self) -> str: return self._session_id diff --git a/python/pyspark/sql/session.py b/python/pyspark/sql/session.py index c36260b9d13ea..32f8f7659b66a 100644 --- a/python/pyspark/sql/session.py +++ b/python/pyspark/sql/session.py @@ -523,7 +523,27 @@ def getOrCreate(self) -> "SparkSession": messageParameters={}, ) - if url.startswith("local") or ( + reuse_local = str( + opts.get( + "spark.local.connect.reuse", + os.environ.get("SPARK_LOCAL_CONNECT_REUSE", ""), + ) + ).lower() in ("1", "true") + + if url.startswith("local") and reuse_local: + # Opt-in: reconnect to a persistent local Connect server (starting + # one on the first run) instead of booting a fresh in-process server + # every process. See `_reuse_or_start_local_connect_server`. + url = RemoteSparkSession._reuse_or_start_local_connect_server( + url, opts + ) + for k in ( + "spark.local.connect.reuse", + "spark.local.connect.server.port", + "spark.local.connect.server.idleTimeout", + ): + opts.pop(k, None) + elif url.startswith("local") or ( is_api_mode_connect and not url.startswith("sc://") ): os.environ["SPARK_LOCAL_REMOTE"] = "1" diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server.py b/python/pyspark/sql/tests/connect/test_connect_local_server.py new file mode 100644 index 0000000000000..94e0ac7a08804 --- /dev/null +++ b/python/pyspark/sql/tests/connect/test_connect_local_server.py @@ -0,0 +1,199 @@ +# +# 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 json +import os +import socket +import tempfile +import time +import unittest + +from pyspark.util import is_remote_only +from pyspark.testing.connectutils import should_test_connect, connect_requirement_message + +if should_test_connect: + from pyspark.sql.connect.session import SparkSession as RemoteSparkSession + from pyspark.version import __version__ + + +@unittest.skipIf( + not should_test_connect or is_remote_only(), + connect_requirement_message or "Requires JVM access to start a local Connect server", +) +class LocalConnectServerReuseTests(unittest.TestCase): + """Tests for the opt-in persistent local Spark Connect server (SPARK_LOCAL_CONNECT_REUSE).""" + + def setUp(self) -> None: + # Point discovery at a throwaway path and remember the env we override, so each test starts + # from a clean slate and the real ~/.spark/connect-local.json is never touched. + self._tmpdir = tempfile.mkdtemp() + self._discovery = os.path.join(self._tmpdir, "connect-local.json") + self._saved_env = { + k: os.environ.get(k) + for k in ("SPARK_LOCAL_CONNECT_DISCOVERY", "SPARK_CONNECT_AUTHENTICATE_TOKEN") + } + os.environ["SPARK_LOCAL_CONNECT_DISCOVERY"] = self._discovery + + def tearDown(self) -> None: + try: + # Only stop a real, separately-spawned server. The discovery-logic unit tests fabricate + # discovery files that point at this very process, which must never be signalled. + disc = RemoteSparkSession._read_local_connect_discovery() + if disc is not None and disc.get("pid") != os.getpid(): + RemoteSparkSession._stop_local_connect_server() + finally: + for k, v in self._saved_env.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + try: + os.remove(self._discovery) + except OSError: + pass + os.rmdir(self._tmpdir) + + # -- discovery / reuse-decision logic (no real server) ---------------------------------------- + + def test_discovery_path_honors_override(self) -> None: + self.assertEqual(RemoteSparkSession._local_connect_discovery_path(), self._discovery) + os.environ.pop("SPARK_LOCAL_CONNECT_DISCOVERY") + self.assertTrue( + RemoteSparkSession._local_connect_discovery_path().endswith( + os.path.join(".spark", "connect-local.json") + ) + ) + + def test_read_discovery_missing_or_malformed(self) -> None: + self.assertIsNone(RemoteSparkSession._read_local_connect_discovery()) + with open(self._discovery, "w") as f: + f.write("not json") + self.assertIsNone(RemoteSparkSession._read_local_connect_discovery()) + with open(self._discovery, "w") as f: + json.dump({"host": "localhost"}, f) # missing required keys + self.assertIsNone(RemoteSparkSession._read_local_connect_discovery()) + + def _write_discovery(self, **overrides) -> dict: + disc = { + "host": "localhost", + "port": 0, + "token": "t", + "pid": os.getpid(), + "spark_version": __version__, + } + disc.update(overrides) + with open(self._discovery, "w") as f: + json.dump(disc, f) + return disc + + def test_not_reusable_on_version_mismatch(self) -> None: + disc = self._write_discovery(spark_version="0.0.0-not-this-build") + self.assertFalse(RemoteSparkSession._local_connect_server_is_reusable(disc)) + + def test_not_reusable_on_dead_pid(self) -> None: + # PID 2**31 - 1 is effectively guaranteed not to exist. + disc = self._write_discovery(pid=2**31 - 1, port=1) + self.assertFalse(RemoteSparkSession._local_connect_server_is_reusable(disc)) + + def test_reusable_when_alive_and_listening(self) -> None: + # A live listening socket owned by this (alive) process with a matching version is reusable. + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + listener.bind(("localhost", 0)) + listener.listen(1) + port = listener.getsockname()[1] + disc = self._write_discovery(port=port) + self.assertTrue(RemoteSparkSession._local_connect_server_is_reusable(disc)) + finally: + listener.close() + # Once the socket is closed the port is no longer reachable, so it is not reusable. + self.assertFalse(RemoteSparkSession._local_connect_server_is_reusable(disc)) + + def test_stop_when_no_server_is_safe(self) -> None: + self.assertFalse(RemoteSparkSession._stop_local_connect_server()) + + # -- end-to-end: start a real detached server, reconnect to it, verify isolation -------------- + + def _release(self, session) -> None: + """Close one client session without stopping the shared server.""" + try: + session.client.release_session() + except Exception: + pass + try: + session.client.close() + except Exception: + pass + + def test_start_reuse_and_session_isolation(self) -> None: + # First call starts a detached persistent server and records it in the discovery file. + endpoint = RemoteSparkSession._reuse_or_start_local_connect_server("local[2]", {}) + self.assertTrue(endpoint.startswith("sc://localhost:")) + + disc = RemoteSparkSession._read_local_connect_discovery() + self.assertIsNotNone(disc) + self.assertEqual(disc["spark_version"], __version__) + self.assertEqual(os.environ.get("SPARK_CONNECT_AUTHENTICATE_TOKEN"), disc["token"]) + first_pid = disc["pid"] + + s1 = s2 = None + try: + # A second call reuses the running server: same endpoint, no new process spawned. + endpoint2 = RemoteSparkSession._reuse_or_start_local_connect_server("local[2]", {}) + self.assertEqual(endpoint2, endpoint) + self.assertEqual(RemoteSparkSession._read_local_connect_discovery()["pid"], first_pid) + + # Two independent client connections to the same server run real queries... + s1 = RemoteSparkSession.builder.remote(endpoint).create() + s2 = RemoteSparkSession.builder.remote(endpoint).create() + self.assertEqual(s1.range(5).count(), 5) + self.assertEqual(s2.range(3).count(), 3) + + # ...and session-local state (a temp view) does not leak across connections. + s1.range(1).createOrReplaceTempView("only_in_s1") + self.assertIn("only_in_s1", [t.name for t in s1.catalog.listTables()]) + self.assertNotIn("only_in_s1", [t.name for t in s2.catalog.listTables()]) + finally: + if s1 is not None: + self._release(s1) + if s2 is not None: + self._release(s2) + + # Stopping signals the server and removes the discovery file. + self.assertTrue(RemoteSparkSession._stop_local_connect_server()) + self.assertIsNone(RemoteSparkSession._read_local_connect_discovery()) + # The server should stop accepting connections shortly afterwards. (We check the port rather + # than the pid: the detached server is a child of this long-lived test process, so once it + # exits it lingers as an unreaped zombie for which os.kill(pid, 0) still succeeds.) + _, _, hostport = endpoint.partition("sc://") + host, _, port = hostport.partition(":") + deadline = time.time() + 30 + closed = False + while time.time() < deadline: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(0.5) + if sock.connect_ex((host, int(port))) != 0: + closed = True + break + time.sleep(0.5) + self.assertTrue(closed, "server port {} still open after stop".format(port)) + + +if __name__ == "__main__": + from pyspark.testing import main + + main() From 7976488d10b7af01217592ee1afcbdd0a704d29a Mon Sep 17 00:00:00 2001 From: ericm-db Date: Tue, 30 Jun 2026 18:05:50 +0000 Subject: [PATCH 02/14] [SPARK-57787][CONNECT] Tighten comments, docstrings and docs Trim verbose comments/docstrings, drop redundant noqa and the private-method example from the user docs, and consolidate exception handling. No functional change. --- docs/spark-connect-overview.md | 9 +++---- python/pyspark/sql/connect/local_server.py | 31 +++++++++------------- python/pyspark/sql/connect/session.py | 20 +++++--------- 3 files changed, 23 insertions(+), 37 deletions(-) diff --git a/docs/spark-connect-overview.md b/docs/spark-connect-overview.md index f0ee4f54db2de..759df7cd07e64 100644 --- a/docs/spark-connect-overview.md +++ b/docs/spark-connect-overview.md @@ -332,12 +332,9 @@ This is **off by default**; nothing changes unless you opt in. A few details: with the host, port, pid and Spark version, to `~/.spark/connect-local.json` (mode `0600`). Set `SPARK_LOCAL_CONNECT_DISCOVERY` to relocate that file. Reuse is refused (and a fresh server started) if the recorded process is gone, the port is closed, or the Spark version differs. -- The server self-terminates after it has been idle for `spark.local.connect.server.idleTimeout` - seconds (default `3600`; set `0` to disable). To stop it explicitly: - -{% highlight bash %} -python -c "from pyspark.sql.connect.session import SparkSession; SparkSession._stop_local_connect_server()" -{% endhighlight %} +- The server self-terminates once it has been idle for `spark.local.connect.server.idleTimeout` + seconds (default `3600`; set `0` to disable). To stop it sooner, terminate the `pid` recorded in + the discovery file. ## Use Spark Connect in standalone applications diff --git a/python/pyspark/sql/connect/local_server.py b/python/pyspark/sql/connect/local_server.py index a5d5298eb3662..f7eccf8032e29 100644 --- a/python/pyspark/sql/connect/local_server.py +++ b/python/pyspark/sql/connect/local_server.py @@ -30,29 +30,25 @@ 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 (not ``python -m pyspark.sql.connect.local_server``) on purpose: it only -needs a classic PySpark install plus the Spark Connect server jar on the classpath -- exactly like -``sbin/start-connect-server.sh`` -- and must not require the Spark Connect *client* dependencies -(grpc, etc.). Importing the ``pyspark.sql.connect`` package would trigger that client-dependency -check, so this file imports only the classic ``pyspark.sql`` API. +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 -# This module is launched by file path (see SparkSession._start_persistent_local_connect_server), -# which puts its own directory -- pyspark/sql/connect -- at the front of sys.path. That directory -# holds modules whose names collide with the standard library (e.g. `types`, `logging`), so leaving -# it there shadows the stdlib and breaks ordinary imports. Drop it before importing anything else; -# pyspark itself stays importable via the remaining sys.path entries (PYTHONPATH / installation). -# `import sys` is a built-in and cannot be shadowed, so it is safe to run first. +# 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 # noqa: E402 -import json # noqa: E402 -import os # noqa: E402 -import signal # noqa: E402 -import time # noqa: E402 -from typing import Any # noqa: E402 +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: @@ -156,7 +152,6 @@ def main() -> None: bound_port = _bound_port(spark, args.port) _write_discovery(args.discovery, "localhost", bound_port, args.token, __version__) - # Printed to the (normally discarded) child stdout; useful when launched with stdout attached. print( "SPARK-CONNECT-LOCAL-SERVER READY port={} pid={}".format(bound_port, os.getpid()), flush=True, diff --git a/python/pyspark/sql/connect/session.py b/python/pyspark/sql/connect/session.py index 2542776d0df0c..440ef937a4d85 100644 --- a/python/pyspark/sql/connect/session.py +++ b/python/pyspark/sql/connect/session.py @@ -1249,15 +1249,11 @@ def _start_connect_server(master: str, opts: Dict[str, Any]) -> None: messageParameters={}, ) - # --------------------------------------------------------------------------------------------- - # Opt-in reuse of a persistent local Spark Connect server. - # - # By default ``SparkSession.builder.remote("local[*]")`` boots a fresh in-process Connect server - # in every process (see ``_start_connect_server`` above), so every ``python script.py`` run - # re-pays the cold start (JVM warmup + SparkContext + server boot). When the opt-in is enabled, - # the first run starts a *detached* server (``connect/local_server.py``) and records it in a - # discovery file; later runs reconnect to it in a fraction of a second instead. - # --------------------------------------------------------------------------------------------- + # 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: @@ -1297,12 +1293,10 @@ def _local_connect_server_is_reusable(disc: Dict[str, Any]) -> bool: return False try: os.kill(int(disc["pid"]), 0) - except ProcessLookupError: - return False - except (TypeError, ValueError): + except (ProcessLookupError, ValueError, TypeError): return False except OSError: - # e.g. PermissionError: the process exists but is not ours to signal -- treat as alive. + # 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) From df7b09a1cd47e6696ef2d77a9dfbf3c1f9042e11 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Wed, 1 Jul 2026 16:56:55 +0000 Subject: [PATCH 03/14] [SPARK-57787][CONNECT] Address review: fix concurrent startup race and seed server confs Fixes two issues raised in review of the persistent local Connect server path: 1. Concurrent first-time startup could fail spuriously. Opted-in processes starting at the same time (or after a version upgrade) would each spawn a server on the fixed port; the loser's JVM failed to bind and raised LOCAL_CONNECT_SERVER_START_FAILED even though a usable server was now up. Serialize start-up with a discovery-file lock, reconnect to the winner if our own daemon exits, and fall back to an ephemeral port when the configured one is already held. 2. The daemon dropped most builder options on first startup. Forward the caller's start-up confs (warehouse dir, jars/packages, catalog, app name, etc.) to the server via a JSON conf file so first-run behavior matches the in-process path. A later run reconnecting to an already-warm JVM still cannot change its static confs, as documented. --- python/pyspark/sql/connect/local_server.py | 15 +- python/pyspark/sql/connect/session.py | 209 ++++++++++++++---- .../connect/test_connect_local_server.py | 81 ++++++- 3 files changed, 259 insertions(+), 46 deletions(-) diff --git a/python/pyspark/sql/connect/local_server.py b/python/pyspark/sql/connect/local_server.py index f7eccf8032e29..7d0108b5625ec 100644 --- a/python/pyspark/sql/connect/local_server.py +++ b/python/pyspark/sql/connect/local_server.py @@ -119,6 +119,11 @@ def main() -> None: 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, @@ -137,9 +142,15 @@ def main() -> None: 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 = ( - SparkSession.builder.master(args.master) - .config("spark.plugins", "org.apache.spark.sql.connect.SparkConnectPlugin") + 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. diff --git a/python/pyspark/sql/connect/session.py b/python/pyspark/sql/connect/session.py index 440ef937a4d85..22152bc7f86c8 100644 --- a/python/pyspark/sql/connect/session.py +++ b/python/pyspark/sql/connect/session.py @@ -1305,39 +1305,156 @@ def _local_connect_server_is_reusable(disc: Dict[str, Any]) -> bool: return True @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. + def _reuse_from_discovery() -> Optional[str]: + """Return an endpoint for the recorded server if it is reusable, else ``None``. - Returns the ``sc://host:port`` endpoint to connect to and sets - ``SPARK_CONNECT_AUTHENTICATE_TOKEN`` so the client authenticates against the server. 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. + 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 SparkSession._start_persistent_local_connect_server(master, opts) + 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``, or ``None`` when file + locking is unavailable (e.g. Windows, which has no ``fcntl``); there we rely on the + reconnect-the-winner fallback in ``_start_persistent_local_connect_server`` instead. + """ + 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 so concurrent opted-in + # processes (parallel workers, an IDE spawning scripts) do not each spawn a server and + # collide on the port; the winner writes the discovery file and the others reuse 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 _start_persistent_local_connect_server(master: str, opts: Dict[str, Any]) -> str: - """Launch a detached persistent local Connect server and wait until it is reachable.""" + """Launch a detached persistent local Connect server and wait until it is reachable. + + Callers must hold the start-up lock (see ``_reuse_or_start_local_connect_server``). If the + server cannot be started but another process has meanwhile published a reusable one, this + reconnects to it rather than failing. + """ import socket import subprocess + import tempfile discovery_path = SparkSession._local_connect_discovery_path() token = opts.get("spark.connect.authenticate.token") or str(uuid.uuid4()) - # Tests use an ephemeral port (0) so they can run in parallel; the server reports the actual - # bound port back through the discovery file. + # Choose the port. Tests use an ephemeral port (0) so they can run in parallel. Otherwise we + # honor the configured/default port, but fall back to an ephemeral one if it is already + # taken -- e.g. by a stale server we just rejected on version mismatch -- so a fresh server + # can still start instead of failing to bind. if "SPARK_TESTING" in os.environ: port = 0 else: port = int( opts.get("spark.local.connect.server.port", DefaultChannelBuilder.default_port()) ) + if port != 0 and not SparkSession._local_port_available(port): + port = 0 idle_timeout = opts.get("spark.local.connect.server.idleTimeout", "3600") + # Seed the server with the caller's start-up confs (warehouse dir, jars, catalog, etc.) so + # first-run behavior matches the in-process path. Passed as a JSON file since confs are + # arbitrary key/values. Written under the discovery dir with 0600 perms as it may hold + # sensitive values, and removed once the daemon has started. + conf_file = None + seed_conf = SparkSession._local_connect_server_conf(opts) + if seed_conf: + parent = os.path.dirname(discovery_path) + if parent and not os.path.isdir(parent): + os.makedirs(parent, exist_ok=True) + fd, conf_file = tempfile.mkstemp(prefix="connect-local-conf-", dir=parent or None) + with os.fdopen(fd, "w") as f: + json.dump(seed_conf, f) + os.chmod(conf_file, 0o600) + daemon = os.path.join(os.path.dirname(os.path.abspath(__file__)), "local_server.py") cmd = [ sys.executable, @@ -1353,6 +1470,8 @@ def _start_persistent_local_connect_server(master: str, opts: Dict[str, Any]) -> "--idle-timeout", str(idle_timeout), ] + if conf_file is not None: + cmd += ["--conf-file", conf_file] # Launch detached so the server outlives this client process. env = dict(os.environ) @@ -1372,36 +1491,48 @@ def _start_persistent_local_connect_server(master: str, opts: Dict[str, Any]) -> popen_kwargs["creationflags"] = detached | new_group proc = subprocess.Popen(cmd, **popen_kwargs) - # Wait for the server to write the discovery file (which records the actual bound port) and - # start accepting connections. - deadline = time.time() + 120 - while time.time() < deadline: - exit_code = proc.poll() - if exit_code is not None: - raise PySparkRuntimeError( - errorClass="LOCAL_CONNECT_SERVER_START_FAILED", - messageParameters={ - "reason": "the server process exited with code {}".format(exit_code) - }, - ) - disc = SparkSession._read_local_connect_discovery() - if disc is not None and disc.get("pid") == proc.pid and disc.get("token") == token: - 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: - os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = token - return "sc://{}:{}".format(disc["host"], disc["port"]) - time.sleep(0.25) - - # Timed out: do not leave the detached process orphaned. try: - proc.terminate() - except OSError: - pass - raise PySparkRuntimeError( - errorClass="LOCAL_CONNECT_SERVER_START_FAILED", - messageParameters={"reason": "the server did not become ready within 120 seconds"}, - ) + # Wait for the server to write the discovery file (which records the actual bound port) + # and start accepting connections. + deadline = time.time() + 120 + while time.time() < deadline: + exit_code = proc.poll() + if exit_code is not None: + # Our daemon died. Another process may have published a usable server in the + # meantime (e.g. a port race), so prefer reconnecting to it over failing. + endpoint = SparkSession._reuse_from_discovery() + if endpoint is not None: + return endpoint + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={ + "reason": "the server process exited with code {}".format(exit_code) + }, + ) + disc = SparkSession._read_local_connect_discovery() + if disc is not None and disc.get("pid") == proc.pid and disc.get("token") == token: + 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: + os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = token + return "sc://{}:{}".format(disc["host"], disc["port"]) + time.sleep(0.25) + + # Timed out: do not leave the detached process orphaned. + try: + proc.terminate() + except OSError: + pass + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={"reason": "the server did not become ready within 120 seconds"}, + ) + finally: + if conf_file is not None: + try: + os.remove(conf_file) + except OSError: + pass @staticmethod def _stop_local_connect_server() -> bool: diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server.py b/python/pyspark/sql/tests/connect/test_connect_local_server.py index 94e0ac7a08804..3a08160218bbe 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server.py @@ -17,6 +17,7 @@ import json import os +import shutil import socket import tempfile import time @@ -61,11 +62,9 @@ def tearDown(self) -> None: os.environ.pop(k, None) else: os.environ[k] = v - try: - os.remove(self._discovery) - except OSError: - pass - os.rmdir(self._tmpdir) + # Remove the whole scratch dir: besides the discovery file it may hold a .lock file, + # seed-conf temp files, and a seeded warehouse directory. + shutil.rmtree(self._tmpdir, ignore_errors=True) # -- discovery / reuse-decision logic (no real server) ---------------------------------------- @@ -126,6 +125,57 @@ def test_reusable_when_alive_and_listening(self) -> None: def test_stop_when_no_server_is_safe(self) -> None: self.assertFalse(RemoteSparkSession._stop_local_connect_server()) + def test_reuse_from_discovery_none_when_absent(self) -> None: + self.assertIsNone(RemoteSparkSession._reuse_from_discovery()) + + def test_local_port_available(self) -> None: + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + listener.bind(("localhost", 0)) + listener.listen(1) + taken = listener.getsockname()[1] + self.assertFalse(RemoteSparkSession._local_port_available(taken)) + finally: + listener.close() + # The port is free again once the listener is closed. + self.assertTrue(RemoteSparkSession._local_port_available(taken)) + + def test_server_conf_seeds_user_confs_and_drops_control_keys(self) -> None: + """_local_connect_server_conf keeps user startup confs but not keys the daemon controls.""" + opts = { + "spark.sql.warehouse.dir": "/tmp/wh", + "spark.jars.packages": "org.example:lib:1.0", + "spark.remote": "local[*]", + "spark.master": "local[*]", + "spark.connect.authenticate.token": "secret", + "spark.connect.grpc.binding.port": "15002", + "spark.local.connect.reuse": "true", + "spark.local.connect.server.port": "15002", + "spark.local.connect.server.idleTimeout": "60", + } + conf = RemoteSparkSession._local_connect_server_conf(opts) + self.assertEqual(conf.get("spark.sql.warehouse.dir"), "/tmp/wh") + self.assertEqual(conf.get("spark.jars.packages"), "org.example:lib:1.0") + for dropped in ( + "spark.remote", + "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", + ): + self.assertNotIn(dropped, conf) + + def test_start_lock_roundtrip(self) -> None: + """Acquiring and releasing the start-up lock creates the lock file and does not error.""" + fd = RemoteSparkSession._acquire_local_connect_start_lock() + try: + if fd is not None: # None only where fcntl is unavailable (e.g. Windows) + self.assertTrue(os.path.exists(self._discovery + ".lock")) + finally: + RemoteSparkSession._release_local_connect_start_lock(fd) + # -- end-to-end: start a real detached server, reconnect to it, verify isolation -------------- def _release(self, session) -> None: @@ -192,6 +242,27 @@ def test_start_reuse_and_session_isolation(self) -> None: time.sleep(0.5) self.assertTrue(closed, "server port {} still open after stop".format(port)) + def test_start_seeds_static_conf_on_the_server(self) -> None: + """A start-up conf passed by the first caller reaches the daemon's SparkConf. + + Uses a static conf (``spark.sql.warehouse.dir``) that the per-session ``_apply_options`` + path cannot set on an already-running JVM, so observing it on the server proves the seed + conf was forwarded rather than applied client-side afterwards. + """ + warehouse = os.path.join(self._tmpdir, "seeded-wh") + opts = {"spark.sql.warehouse.dir": warehouse} + endpoint = RemoteSparkSession._reuse_or_start_local_connect_server("local[2]", opts) + spark = None + try: + spark = RemoteSparkSession.builder.remote(endpoint).create() + # Spark normalizes the warehouse dir to a file: URI; the path still identifies our dir, + # which proves the seed conf reached the daemon (a per-session apply cannot set it). + self.assertTrue(spark.conf.get("spark.sql.warehouse.dir").endswith(warehouse)) + finally: + if spark is not None: + self._release(spark) + RemoteSparkSession._stop_local_connect_server() + if __name__ == "__main__": from pyspark.testing import main From bd938e71c7d517096faca4f1a6454eb450332211 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Wed, 1 Jul 2026 20:02:06 +0000 Subject: [PATCH 04/14] [SPARK-57787][CONNECT] Address review: reap the JVM on start-up timeout On the start-up timeout path, terminating only the detached daemon process left its child JVM orphaned when the timeout fired before the daemon had installed its own signal handling (the slow- startup case that triggers the timeout). Signal the whole process group -- the daemon is a session leader and the JVM stays in its group -- and escalate to SIGKILL if a graceful stop does not take, so the JVM is reaped instead of leaking. --- python/pyspark/sql/connect/session.py | 34 +++++++++++++++++++++++---- 1 file changed, 29 insertions(+), 5 deletions(-) diff --git a/python/pyspark/sql/connect/session.py b/python/pyspark/sql/connect/session.py index 22152bc7f86c8..f38282b408fd3 100644 --- a/python/pyspark/sql/connect/session.py +++ b/python/pyspark/sql/connect/session.py @@ -1411,6 +1411,30 @@ def _local_connect_server_conf(opts: Dict[str, Any]) -> Dict[str, Any]: conf.pop(k, None) return conf + @staticmethod + def _terminate_local_connect_server(proc: Any) -> None: + """Terminate a daemon started by ``_start_persistent_local_connect_server`` and its JVM. + + The daemon is a POSIX session leader (``start_new_session=True``) and its child JVM stays in + that process group, so signalling the group reaps both -- important when the timeout fires + before the daemon has wired up its own signal handling. Escalates to SIGKILL if a graceful + stop does not take. On non-POSIX platforms it falls back to terminating just the process. + """ + try: + if os.name == "posix": + os.killpg(os.getpgid(proc.pid), signal.SIGTERM) + else: + proc.terminate() + try: + proc.wait(timeout=10) + except Exception: + if os.name == "posix": + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + else: + proc.kill() + except OSError: + pass + @staticmethod def _start_persistent_local_connect_server(master: str, opts: Dict[str, Any]) -> str: """Launch a detached persistent local Connect server and wait until it is reachable. @@ -1518,11 +1542,11 @@ def _start_persistent_local_connect_server(master: str, opts: Dict[str, Any]) -> return "sc://{}:{}".format(disc["host"], disc["port"]) time.sleep(0.25) - # Timed out: do not leave the detached process orphaned. - try: - proc.terminate() - except OSError: - pass + # Timed out. The daemon may still be inside getOrCreate() -- i.e. before it has wired up + # its own SIGTERM handler -- and it has already spawned a child JVM. Signal the whole + # process group (the daemon is a session leader via start_new_session) and escalate to + # SIGKILL, so the JVM is reaped rather than orphaned. + SparkSession._terminate_local_connect_server(proc) raise PySparkRuntimeError( errorClass="LOCAL_CONNECT_SERVER_START_FAILED", messageParameters={"reason": "the server did not become ready within 120 seconds"}, From 321128d20b07acf634559bd0e129d458b8a8b957 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Wed, 1 Jul 2026 21:55:41 +0000 Subject: [PATCH 05/14] [SPARK-57787][CONNECT] Add test for reaping the JVM on start-up timeout Adds a POSIX-only test that starts the detached daemon, waits until it has spawned its child JVM, then calls _terminate_local_connect_server and asserts the whole process group is gone -- covering the orphaned-JVM case the timeout-path fix addresses. Also waits for the server port to close in tearDown so a stopped server's JVM cannot linger into the next test. --- .../connect/test_connect_local_server.py | 105 ++++++++++++++++-- 1 file changed, 94 insertions(+), 11 deletions(-) diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server.py b/python/pyspark/sql/tests/connect/test_connect_local_server.py index 3a08160218bbe..c86b15db29cf7 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server.py @@ -18,7 +18,9 @@ import json import os import shutil +import signal import socket +import sys import tempfile import time import unittest @@ -55,7 +57,11 @@ def tearDown(self) -> None: # discovery files that point at this very process, which must never be signalled. disc = RemoteSparkSession._read_local_connect_discovery() if disc is not None and disc.get("pid") != os.getpid(): + port = int(disc["port"]) RemoteSparkSession._stop_local_connect_server() + # _stop_local_connect_server only signals the daemon and returns; wait for the JVM + # to actually release the port so the next test starts from a clean slate. + self._wait_port_closed(disc["host"], port) finally: for k, v in self._saved_env.items(): if v is None: @@ -189,6 +195,17 @@ def _release(self, session) -> None: except Exception: pass + def _wait_port_closed(self, host, port, timeout=30) -> bool: + """Wait for ``host:port`` to stop accepting connections; return True if it closed.""" + deadline = time.time() + timeout + while time.time() < deadline: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(0.5) + if sock.connect_ex((host, int(port))) != 0: + return True + time.sleep(0.5) + return False + def test_start_reuse_and_session_isolation(self) -> None: # First call starts a detached persistent server and records it in the discovery file. endpoint = RemoteSparkSession._reuse_or_start_local_connect_server("local[2]", {}) @@ -231,16 +248,9 @@ def test_start_reuse_and_session_isolation(self) -> None: # exits it lingers as an unreaped zombie for which os.kill(pid, 0) still succeeds.) _, _, hostport = endpoint.partition("sc://") host, _, port = hostport.partition(":") - deadline = time.time() + 30 - closed = False - while time.time() < deadline: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.settimeout(0.5) - if sock.connect_ex((host, int(port))) != 0: - closed = True - break - time.sleep(0.5) - self.assertTrue(closed, "server port {} still open after stop".format(port)) + self.assertTrue( + self._wait_port_closed(host, port), "server port {} still open after stop".format(port) + ) def test_start_seeds_static_conf_on_the_server(self) -> None: """A start-up conf passed by the first caller reaches the daemon's SparkConf. @@ -261,7 +271,80 @@ def test_start_seeds_static_conf_on_the_server(self) -> None: finally: if spark is not None: self._release(spark) - RemoteSparkSession._stop_local_connect_server() + # tearDown stops the server and waits for the port to close. + + @unittest.skipUnless(os.name == "posix", "process-group reaping is POSIX-only") + def test_terminate_reaps_daemon_and_jvm(self) -> None: + """_terminate_local_connect_server kills the daemon *and* its child JVM. + + Reproduces the start-up-timeout orphan case: terminating only the daemon leader would leak + the JVM it spawned. The helper signals the whole process group, so nothing in the group + survives -- checked here after the daemon has actually launched its JVM. + """ + import subprocess + + if shutil.which("pgrep") is None: + self.skipTest("pgrep is needed to detect the child JVM") + + import pyspark.sql.connect.session as session_mod + + daemon = os.path.join(os.path.dirname(session_mod.__file__), "local_server.py") + cmd = [ + sys.executable, daemon, + "--master", "local[2]", + "--port", "0", # ephemeral, so a stray server elsewhere cannot interfere + "--token", "reap-test", + "--discovery", self._discovery, + "--idle-timeout", "3600", + ] + env = dict(os.environ) + for var in ("SPARK_REMOTE", "SPARK_LOCAL_REMOTE", "SPARK_CONNECT_MODE_ENABLED"): + env.pop(var, None) + proc = subprocess.Popen( + cmd, env=env, + stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + start_new_session=True, # session/group leader, matching the production launch + ) + + def child_pids(): + try: + out = subprocess.check_output(["pgrep", "-P", str(proc.pid)]) + except subprocess.CalledProcessError: + return [] + return [int(p) for p in out.split()] + + pgid = os.getpgid(proc.pid) + try: + # Wait until the daemon has spawned its child JVM -- the process the fix must also reap. + deadline = time.time() + 90 + kids = [] + while time.time() < deadline: + kids = child_pids() + if kids: + break + if proc.poll() is not None: + self.fail("daemon exited (code {}) before launching a JVM".format(proc.poll())) + time.sleep(0.5) + self.assertTrue(kids, "daemon never launched a JVM; cannot exercise the orphan case") + + RemoteSparkSession._terminate_local_connect_server(proc) + + # killpg(sig 0) raises ProcessLookupError once no process in the group remains. + deadline = time.time() + 30 + while time.time() < deadline: + try: + os.killpg(pgid, 0) + except ProcessLookupError: + break + time.sleep(0.5) + with self.assertRaises(ProcessLookupError): + os.killpg(pgid, 0) + finally: + # Belt and suspenders: never leak the group if an assertion above failed early. + try: + os.killpg(pgid, signal.SIGKILL) + except OSError: + pass if __name__ == "__main__": From c99cf02e35082562bdf9599c7e9d36c92ec18b7a Mon Sep 17 00:00:00 2001 From: ericm-db Date: Tue, 7 Jul 2026 22:54:49 +0000 Subject: [PATCH 06/14] [SPARK-57787][CONNECT] Address local Connect reuse review comments --- docs/spark-connect-overview.md | 13 +- python/pyspark/sql/connect/session.py | 46 +++--- .../connect/test_connect_local_server.py | 137 +++++++++++++++--- 3 files changed, 143 insertions(+), 53 deletions(-) diff --git a/docs/spark-connect-overview.md b/docs/spark-connect-overview.md index 759df7cd07e64..be0d2ee31376b 100644 --- a/docs/spark-connect-overview.md +++ b/docs/spark-connect-overview.md @@ -331,10 +331,15 @@ This is **off by default**; nothing changes unless you opt in. A few details: - The server listens on port `15002` by default and authenticates with a token written, together with the host, port, pid and Spark version, to `~/.spark/connect-local.json` (mode `0600`). Set `SPARK_LOCAL_CONNECT_DISCOVERY` to relocate that file. Reuse is refused (and a fresh server - started) if the recorded process is gone, the port is closed, or the Spark version differs. + started) if the recorded process is gone, the port is closed, or the Spark version differs. On + Unix-like systems, PySpark uses a file lock around first startup. On platforms without `fcntl`, + concurrent startups can race; callers that lose the race reconnect to the server recorded in the + discovery file. - The server self-terminates once it has been idle for `spark.local.connect.server.idleTimeout` - seconds (default `3600`; set `0` to disable). To stop it sooner, terminate the `pid` recorded in - the discovery file. + seconds (default `3600`; set `0` to disable). To stop it sooner, call + `SparkSession._stop_local_connect_server()` or terminate the `pid` recorded in the discovery file. + If an old discovery file is rejected after a Spark upgrade, stop the old recorded `pid` if you do + not want to wait for the idle timeout. ## Use Spark Connect in standalone applications @@ -430,7 +435,7 @@ one may implement their own class extending `ClassFinder` for customized search
For more information on application development with Spark Connect as well as extending Spark Connect -with custom functionality, see [Application Development with Spark Connect](app-dev-spark-connect.html). +with custom functionality, see [Application Development with Spark Connect](app-dev-spark-connect.html). # Client application authentication While Spark Connect does not have built-in authentication, it is designed to diff --git a/python/pyspark/sql/connect/session.py b/python/pyspark/sql/connect/session.py index f38282b408fd3..51627ab4a4343 100644 --- a/python/pyspark/sql/connect/session.py +++ b/python/pyspark/sql/connect/session.py @@ -1321,9 +1321,9 @@ def _reuse_from_discovery() -> Optional[str]: 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``, or ``None`` when file - locking is unavailable (e.g. Windows, which has no ``fcntl``); there we rely on the - reconnect-the-winner fallback in ``_start_persistent_local_connect_server`` instead. + 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 @@ -1360,9 +1360,9 @@ def _reuse_or_start_local_connect_server(master: str, opts: Dict[str, Any]) -> s endpoint = SparkSession._reuse_from_discovery() if endpoint is not None: return endpoint - # No reusable server yet. Serialize start-up across processes so concurrent opted-in - # processes (parallel workers, an IDE spawning scripts) do not each spawn a server and - # collide on the port; the winner writes the discovery file and the others reuse it. + # 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() @@ -1412,28 +1412,28 @@ def _local_connect_server_conf(opts: Dict[str, Any]) -> Dict[str, Any]: return conf @staticmethod - def _terminate_local_connect_server(proc: Any) -> None: - """Terminate a daemon started by ``_start_persistent_local_connect_server`` and its JVM. - - The daemon is a POSIX session leader (``start_new_session=True``) and its child JVM stays in - that process group, so signalling the group reaps both -- important when the timeout fires - before the daemon has wired up its own signal handling. Escalates to SIGKILL if a graceful - stop does not take. On non-POSIX platforms it falls back to terminating just the process. - """ + def _signal_local_connect_server(pid: int, sig: int) -> bool: + """Signal a detached local Connect daemon, including its JVM on POSIX.""" try: if os.name == "posix": - os.killpg(os.getpgid(proc.pid), signal.SIGTERM) + os.killpg(os.getpgid(pid), sig) else: - proc.terminate() + 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.""" + if SparkSession._signal_local_connect_server(proc.pid, signal.SIGTERM): try: proc.wait(timeout=10) except Exception: if os.name == "posix": - os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + SparkSession._signal_local_connect_server(proc.pid, signal.SIGKILL) else: proc.kill() - except OSError: - pass @staticmethod def _start_persistent_local_connect_server(master: str, opts: Dict[str, Any]) -> str: @@ -1576,14 +1576,8 @@ def _stop_local_connect_server() -> bool: pid = int(disc["pid"]) except (TypeError, ValueError, KeyError): pid = None - # Never signal the current process: a discovery file should only ever point at the - # detached server, never at the client that is reading it. if pid is not None and pid != os.getpid(): - try: - os.kill(pid, signal.SIGTERM) - stopped = True - except OSError: - pass + stopped = SparkSession._signal_local_connect_server(pid, signal.SIGTERM) try: os.remove(path) except OSError: diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server.py b/python/pyspark/sql/tests/connect/test_connect_local_server.py index c86b15db29cf7..3642473ea569a 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server.py @@ -20,8 +20,10 @@ import shutil import signal import socket +import subprocess import sys import tempfile +import textwrap import time import unittest @@ -29,6 +31,7 @@ from pyspark.testing.connectutils import should_test_connect, connect_requirement_message if should_test_connect: + from pyspark.sql import SparkSession as PySparkSession from pyspark.sql.connect.session import SparkSession as RemoteSparkSession from pyspark.version import __version__ @@ -131,6 +134,23 @@ def test_reusable_when_alive_and_listening(self) -> None: def test_stop_when_no_server_is_safe(self) -> None: self.assertFalse(RemoteSparkSession._stop_local_connect_server()) + def test_stop_signals_recorded_server_group(self) -> None: + self._write_discovery(pid=12345) + calls = [] + old_signal = RemoteSparkSession._signal_local_connect_server + + def fake_signal(pid, sig): + calls.append((pid, sig)) + return True + + try: + RemoteSparkSession._signal_local_connect_server = staticmethod(fake_signal) + self.assertTrue(RemoteSparkSession._stop_local_connect_server()) + self.assertEqual(calls, [(12345, signal.SIGTERM)]) + self.assertFalse(os.path.exists(self._discovery)) + finally: + RemoteSparkSession._signal_local_connect_server = old_signal + def test_reuse_from_discovery_none_when_absent(self) -> None: self.assertIsNone(RemoteSparkSession._reuse_from_discovery()) @@ -206,6 +226,78 @@ def _wait_port_closed(self, host, port, timeout=30) -> bool: time.sleep(0.5) return False + def test_builder_remote_local_uses_reuse_flag(self) -> None: + spark = None + try: + spark = ( + PySparkSession.builder.remote("local[2]") + .config("spark.local.connect.reuse", "true") + .getOrCreate() + ) + self.assertEqual(spark.range(2).count(), 2) + + disc = RemoteSparkSession._read_local_connect_discovery() + self.assertIsNotNone(disc) + self.assertEqual(disc["spark_version"], __version__) + self.assertNotEqual(disc["pid"], os.getpid()) + finally: + if spark is not None: + spark.stop() + + def test_concurrent_startup_reuses_one_server(self) -> None: + script = textwrap.dedent( + """ + import json + import os + + from pyspark.sql import SparkSession + + spark = ( + SparkSession.builder.remote("local[2]") + .config("spark.local.connect.reuse", "true") + .getOrCreate() + ) + try: + count = spark.range(1).count() + with open(os.environ["SPARK_LOCAL_CONNECT_DISCOVERY"], "r") as f: + disc = json.load(f) + print(json.dumps({"count": count, "pid": disc["pid"], "port": disc["port"]})) + finally: + spark.stop() + """ + ) + env = dict(os.environ) + env["SPARK_LOCAL_CONNECT_DISCOVERY"] = self._discovery + env["SPARK_LOCAL_CONNECT_REUSE"] = "1" + + procs = [ + subprocess.Popen( + [sys.executable, "-c", script], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + for _ in range(3) + ] + outputs = [] + try: + for proc in procs: + stdout, stderr = proc.communicate(timeout=180) + self.assertEqual(proc.returncode, 0, stderr) + lines = stdout.strip().splitlines() + self.assertTrue(lines, stderr) + outputs.append(json.loads(lines[-1])) + finally: + for proc in procs: + if proc.poll() is None: + proc.kill() + proc.communicate() + + self.assertEqual({o["count"] for o in outputs}, {1}) + self.assertEqual(len({o["pid"] for o in outputs}), 1) + self.assertEqual(len({o["port"] for o in outputs}), 1) + def test_start_reuse_and_session_isolation(self) -> None: # First call starts a detached persistent server and records it in the discovery file. endpoint = RemoteSparkSession._reuse_or_start_local_connect_server("local[2]", {}) @@ -253,20 +345,14 @@ def test_start_reuse_and_session_isolation(self) -> None: ) def test_start_seeds_static_conf_on_the_server(self) -> None: - """A start-up conf passed by the first caller reaches the daemon's SparkConf. - - Uses a static conf (``spark.sql.warehouse.dir``) that the per-session ``_apply_options`` - path cannot set on an already-running JVM, so observing it on the server proves the seed - conf was forwarded rather than applied client-side afterwards. - """ + """A start-up conf passed by the first caller reaches the daemon's SparkConf.""" warehouse = os.path.join(self._tmpdir, "seeded-wh") opts = {"spark.sql.warehouse.dir": warehouse} endpoint = RemoteSparkSession._reuse_or_start_local_connect_server("local[2]", opts) spark = None try: spark = RemoteSparkSession.builder.remote(endpoint).create() - # Spark normalizes the warehouse dir to a file: URI; the path still identifies our dir, - # which proves the seed conf reached the daemon (a per-session apply cannot set it). + # The per-session apply path cannot set this static conf after the JVM is running. self.assertTrue(spark.conf.get("spark.sql.warehouse.dir").endswith(warehouse)) finally: if spark is not None: @@ -277,12 +363,9 @@ def test_start_seeds_static_conf_on_the_server(self) -> None: def test_terminate_reaps_daemon_and_jvm(self) -> None: """_terminate_local_connect_server kills the daemon *and* its child JVM. - Reproduces the start-up-timeout orphan case: terminating only the daemon leader would leak - the JVM it spawned. The helper signals the whole process group, so nothing in the group - survives -- checked here after the daemon has actually launched its JVM. + The test waits until the daemon has launched the JVM, then checks that the whole process + group exits. """ - import subprocess - if shutil.which("pgrep") is None: self.skipTest("pgrep is needed to detect the child JVM") @@ -290,19 +373,28 @@ def test_terminate_reaps_daemon_and_jvm(self) -> None: daemon = os.path.join(os.path.dirname(session_mod.__file__), "local_server.py") cmd = [ - sys.executable, daemon, - "--master", "local[2]", - "--port", "0", # ephemeral, so a stray server elsewhere cannot interfere - "--token", "reap-test", - "--discovery", self._discovery, - "--idle-timeout", "3600", + sys.executable, + daemon, + "--master", + "local[2]", + "--port", + "0", # ephemeral, so a stray server elsewhere cannot interfere + "--token", + "reap-test", + "--discovery", + self._discovery, + "--idle-timeout", + "3600", ] env = dict(os.environ) for var in ("SPARK_REMOTE", "SPARK_LOCAL_REMOTE", "SPARK_CONNECT_MODE_ENABLED"): env.pop(var, None) proc = subprocess.Popen( - cmd, env=env, - stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + cmd, + env=env, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, start_new_session=True, # session/group leader, matching the production launch ) @@ -315,7 +407,7 @@ def child_pids(): pgid = os.getpgid(proc.pid) try: - # Wait until the daemon has spawned its child JVM -- the process the fix must also reap. + # Wait until the daemon has spawned its child JVM before terminating the group. deadline = time.time() + 90 kids = [] while time.time() < deadline: @@ -340,7 +432,6 @@ def child_pids(): with self.assertRaises(ProcessLookupError): os.killpg(pgid, 0) finally: - # Belt and suspenders: never leak the group if an assertion above failed early. try: os.killpg(pgid, signal.SIGKILL) except OSError: From 957c0f4bbc583ee04fd666d31c7eefff5000cd17 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Tue, 7 Jul 2026 16:38:21 -0700 Subject: [PATCH 07/14] [SPARK-57787][CONNECT] Only group-kill session-leader pids when stopping the local server The stop path signals the pid recorded in the discovery file via killpg. If that pid is stale (e.g. the daemon was SIGKILLed and left its discovery file behind) and has been recycled by an unrelated process, group-killing it could take down the caller's own process group. The daemon is always launched as a session leader, so its group id equals its pid; only signal the group when that holds, and fall back to a plain kill otherwise. Adds a unit test covering both directions (with the signal syscalls mocked so a regression cannot kill the test runner), and a ruff-format fix for a missing blank line after the local_server.py module docstring. Co-authored-by: Isaac --- python/pyspark/sql/connect/local_server.py | 1 + python/pyspark/sql/connect/session.py | 12 ++++-- .../connect/test_connect_local_server.py | 39 +++++++++++++++++++ 3 files changed, 49 insertions(+), 3 deletions(-) diff --git a/python/pyspark/sql/connect/local_server.py b/python/pyspark/sql/connect/local_server.py index 7d0108b5625ec..19a9b74830d9a 100644 --- a/python/pyspark/sql/connect/local_server.py +++ b/python/pyspark/sql/connect/local_server.py @@ -34,6 +34,7 @@ *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 diff --git a/python/pyspark/sql/connect/session.py b/python/pyspark/sql/connect/session.py index 51627ab4a4343..e65f959a6c4f5 100644 --- a/python/pyspark/sql/connect/session.py +++ b/python/pyspark/sql/connect/session.py @@ -1413,10 +1413,16 @@ def _local_connect_server_conf(opts: Dict[str, Any]) -> Dict[str, Any]: @staticmethod def _signal_local_connect_server(pid: int, sig: int) -> bool: - """Signal a detached local Connect daemon, including its JVM on POSIX.""" + """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": - os.killpg(os.getpgid(pid), sig) + if os.name == "posix" and os.getpgid(pid) == pid: + os.killpg(pid, sig) else: os.kill(pid, sig) return True diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server.py b/python/pyspark/sql/tests/connect/test_connect_local_server.py index 3642473ea569a..bdc3d53cda00a 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server.py @@ -151,6 +151,45 @@ def fake_signal(pid, sig): finally: RemoteSparkSession._signal_local_connect_server = old_signal + @unittest.skipUnless(os.name == "posix", "process groups are POSIX-only") + def test_signal_group_kills_only_session_leaders(self) -> None: + """A pid that is not a session leader is signalled alone, never via its group. + + The daemon is always launched as a session leader (its group id equals its pid). A + recorded pid whose group id differs was recycled by an unrelated process -- for example + via a stale discovery file -- and group-killing it could take down the caller's own + process group. + """ + from unittest import mock + + sleeper = [sys.executable, "-c", "import time; time.sleep(60)"] + # Same process group as this test (not a leader): must fall back to a plain kill. + child = subprocess.Popen(sleeper) + try: + self.assertNotEqual(os.getpgid(child.pid), child.pid) + with mock.patch("os.killpg") as killpg, mock.patch("os.kill") as kill: + self.assertTrue( + RemoteSparkSession._signal_local_connect_server(child.pid, signal.SIGTERM) + ) + killpg.assert_not_called() + kill.assert_called_once_with(child.pid, signal.SIGTERM) + finally: + child.kill() + child.communicate() + # A session leader, like the real daemon: the whole group is signalled. + leader = subprocess.Popen(sleeper, start_new_session=True) + try: + self.assertEqual(os.getpgid(leader.pid), leader.pid) + with mock.patch("os.killpg") as killpg, mock.patch("os.kill") as kill: + self.assertTrue( + RemoteSparkSession._signal_local_connect_server(leader.pid, signal.SIGTERM) + ) + killpg.assert_called_once_with(leader.pid, signal.SIGTERM) + kill.assert_not_called() + finally: + leader.kill() + leader.communicate() + def test_reuse_from_discovery_none_when_absent(self) -> None: self.assertIsNone(RemoteSparkSession._reuse_from_discovery()) From 7de16e01434909c29a650c7857d9ac33fa6bd49c Mon Sep 17 00:00:00 2001 From: ericm-db Date: Tue, 7 Jul 2026 22:42:42 -0700 Subject: [PATCH 08/14] [SPARK-57787][CONNECT] Escalate to a group SIGKILL when the JVM outlives the daemon test_terminate_reaps_daemon_and_jvm fails in CI with "ProcessLookupError not raised": _terminate_local_connect_server keyed SIGKILL escalation on the daemon timing out, so when the daemon exits promptly on SIGTERM while its JVM's graceful shutdown lingers (or wedges), nothing ever escalates and the group survives. Terminate now waits for the whole process group to disappear after the daemon settles and group-SIGKILLs any survivors past a 10s grace period. It signals the group id directly since the leader may already be reaped by then, and polls the daemon so a zombie leader does not keep the group looking alive. Adds a deterministic regression test (a session leader that exits on SIGTERM after spawning a SIGTERM-immune child standing in for the slow JVM), and hardens both terminate tests against macOS zombie semantics: killpg on a group holding only a zombie yields EPERM, and kill(pid, 0) succeeds for a zombie until it is reaped. Co-authored-by: Isaac --- python/pyspark/sql/connect/session.py | 41 +++++-- .../connect/test_connect_local_server.py | 102 ++++++++++++++++-- 2 files changed, 125 insertions(+), 18 deletions(-) diff --git a/python/pyspark/sql/connect/session.py b/python/pyspark/sql/connect/session.py index e65f959a6c4f5..126e1ed5652bd 100644 --- a/python/pyspark/sql/connect/session.py +++ b/python/pyspark/sql/connect/session.py @@ -1431,15 +1431,40 @@ def _signal_local_connect_server(pid: int, sig: int) -> bool: @staticmethod def _terminate_local_connect_server(proc: Any) -> None: - """Terminate a daemon started by ``_start_persistent_local_connect_server`` and its JVM.""" - if SparkSession._signal_local_connect_server(proc.pid, signal.SIGTERM): + """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: - proc.wait(timeout=10) - except Exception: - if os.name == "posix": - SparkSession._signal_local_connect_server(proc.pid, signal.SIGKILL) - else: - proc.kill() + os.killpg(proc.pid, 0) + except OSError: + return + time.sleep(0.2) + try: + os.killpg(proc.pid, signal.SIGKILL) + except OSError: + pass + proc.poll() @staticmethod def _start_persistent_local_connect_server(master: str, opts: Dict[str, Any]) -> str: diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server.py b/python/pyspark/sql/tests/connect/test_connect_local_server.py index bdc3d53cda00a..c5231d48d961c 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server.py @@ -398,6 +398,94 @@ def test_start_seeds_static_conf_on_the_server(self) -> None: self._release(spark) # tearDown stops the server and waits for the port to close. + @staticmethod + def _wait_process_group_gone(pgid: int, timeout: float) -> bool: + """Wait until no process in ``pgid`` can be signalled; True if that happened in time. + + ProcessLookupError means the group is empty. PermissionError also counts as gone: on + macOS, a dead group member lingering as a zombie (reparented to launchd) yields EPERM + even though nothing in the group is running. + """ + deadline = time.time() + timeout + while True: + try: + os.killpg(pgid, 0) + except (ProcessLookupError, PermissionError): + return True + if time.time() >= deadline: + return False + time.sleep(0.2) + + @unittest.skipUnless(os.name == "posix", "process-group reaping is POSIX-only") + def test_terminate_escalates_when_group_outlives_daemon(self) -> None: + """SIGKILL escalation covers group members that survive the daemon. + + Reproduces the CI failure mode of test_terminate_reaps_daemon_and_jvm: the daemon exits + on SIGTERM promptly while its JVM is still shutting down (or wedged) past the grace + period. _terminate_local_connect_server must not return while the group has live members. + """ + # A session leader that exits on SIGTERM after spawning a SIGTERM-immune child into its + # process group -- the child stands in for a JVM that outlives the daemon. + leader_prog = textwrap.dedent( + """ + import signal + import subprocess + import sys + import time + + child = subprocess.Popen( + [ + sys.executable, + "-c", + "import signal, time; " + "signal.signal(signal.SIGTERM, signal.SIG_IGN); " + "print('ready', flush=True); " + "time.sleep(600)", + ] + ) + print("pid:%d" % child.pid, flush=True) + signal.signal(signal.SIGTERM, lambda *_: sys.exit(0)) + time.sleep(600) + """ + ) + proc = subprocess.Popen( + [sys.executable, "-c", leader_prog], + stdout=subprocess.PIPE, + start_new_session=True, + ) + pgid = os.getpgid(proc.pid) + try: + # One line comes from the leader (the child's pid) and one from the child itself, + # printed only after it has ignored SIGTERM -- reading both closes the start-up race. + lines = {proc.stdout.readline().strip() for _ in range(2)} + self.assertIn(b"ready", lines) + child_pid = int(next(ln for ln in lines if ln.startswith(b"pid:"))[4:]) + + RemoteSparkSession._terminate_local_connect_server(proc) + + self.assertTrue( + self._wait_process_group_gone(pgid, timeout=10), + "process group survived _terminate_local_connect_server", + ) + # The SIGTERM-immune child must be dead, proving the SIGKILL escalation fired. + # os.kill(pid, 0) still succeeds for a zombie until init/launchd reaps it, so poll + # briefly rather than asserting on the first probe. + deadline = time.time() + 10 + child_gone = False + while time.time() < deadline and not child_gone: + try: + os.kill(child_pid, 0) + time.sleep(0.2) + except OSError: + child_gone = True + self.assertTrue(child_gone, "SIGTERM-immune child survived the SIGKILL escalation") + finally: + proc.stdout.close() + try: + os.killpg(pgid, signal.SIGKILL) + except OSError: + pass + @unittest.skipUnless(os.name == "posix", "process-group reaping is POSIX-only") def test_terminate_reaps_daemon_and_jvm(self) -> None: """_terminate_local_connect_server kills the daemon *and* its child JVM. @@ -460,16 +548,10 @@ def child_pids(): RemoteSparkSession._terminate_local_connect_server(proc) - # killpg(sig 0) raises ProcessLookupError once no process in the group remains. - deadline = time.time() + 30 - while time.time() < deadline: - try: - os.killpg(pgid, 0) - except ProcessLookupError: - break - time.sleep(0.5) - with self.assertRaises(ProcessLookupError): - os.killpg(pgid, 0) + self.assertTrue( + self._wait_process_group_gone(pgid, timeout=30), + "process group survived _terminate_local_connect_server", + ) finally: try: os.killpg(pgid, signal.SIGKILL) From f1c8a11ada6e12908e32db2bdf74ce2c3677d593 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Wed, 8 Jul 2026 09:40:59 -0700 Subject: [PATCH 09/14] [SPARK-57787][CONNECT] Assert on running group members in the terminate tests Both terminate tests fail in CI with "process group survived": the tests checked group death via killpg(pgid, 0), but a zombie keeps its group signalable, and in CI containers whose pid 1 never reaps, the daemon's JVM reparents on the daemon's death and lingers as a zombie indefinitely -- so the check never turns negative even though nothing is running. (macOS launchd reaps promptly, which is why the tests pass locally.) The tests now assert that no *running* (non-zombie) process remains in the group, using ps to inspect member states. Co-authored-by: Isaac --- .../connect/test_connect_local_server.py | 48 +++++++++++-------- 1 file changed, 28 insertions(+), 20 deletions(-) diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server.py b/python/pyspark/sql/tests/connect/test_connect_local_server.py index c5231d48d961c..a43a0388df254 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server.py @@ -399,18 +399,29 @@ def test_start_seeds_static_conf_on_the_server(self) -> None: # tearDown stops the server and waits for the port to close. @staticmethod - def _wait_process_group_gone(pgid: int, timeout: float) -> bool: - """Wait until no process in ``pgid`` can be signalled; True if that happened in time. + def _running_group_members(pgid: int) -> list: + """Pids in process group ``pgid`` that are still running, excluding zombies. - ProcessLookupError means the group is empty. PermissionError also counts as gone: on - macOS, a dead group member lingering as a zombie (reparented to launchd) yields EPERM - even though nothing in the group is running. + Zombies count as gone: they are already dead, so nothing can leak. Signal-based checks + like ``killpg(pgid, 0)`` cannot be used instead, because a zombie keeps its group + signalable, and orphans reparented to a pid 1 that never reaps (e.g. minimal CI + containers) stay zombies indefinitely. """ + out = subprocess.run( + ["ps", "-A", "-o", "pid=,pgid=,stat="], capture_output=True, text=True + ).stdout + members = [] + for line in out.splitlines(): + parts = line.split() + if len(parts) >= 3 and parts[1] == str(pgid) and not parts[2].startswith("Z"): + members.append(int(parts[0])) + return members + + def _wait_process_group_gone(self, pgid: int, timeout: float) -> bool: + """Wait until nothing in ``pgid`` is still running; True if that happened in time.""" deadline = time.time() + timeout while True: - try: - os.killpg(pgid, 0) - except (ProcessLookupError, PermissionError): + if not self._running_group_members(pgid): return True if time.time() >= deadline: return False @@ -424,6 +435,9 @@ def test_terminate_escalates_when_group_outlives_daemon(self) -> None: on SIGTERM promptly while its JVM is still shutting down (or wedged) past the grace period. _terminate_local_connect_server must not return while the group has live members. """ + if shutil.which("ps") is None: + self.skipTest("ps is needed to inspect process group members") + # A session leader that exits on SIGTERM after spawning a SIGTERM-immune child into its # process group -- the child stands in for a JVM that outlives the daemon. leader_prog = textwrap.dedent( @@ -467,18 +481,12 @@ def test_terminate_escalates_when_group_outlives_daemon(self) -> None: self._wait_process_group_gone(pgid, timeout=10), "process group survived _terminate_local_connect_server", ) - # The SIGTERM-immune child must be dead, proving the SIGKILL escalation fired. - # os.kill(pid, 0) still succeeds for a zombie until init/launchd reaps it, so poll - # briefly rather than asserting on the first probe. - deadline = time.time() + 10 - child_gone = False - while time.time() < deadline and not child_gone: - try: - os.kill(child_pid, 0) - time.sleep(0.2) - except OSError: - child_gone = True - self.assertTrue(child_gone, "SIGTERM-immune child survived the SIGKILL escalation") + # The SIGTERM-immune child no longer runs, proving the SIGKILL escalation fired. + self.assertNotIn( + child_pid, + self._running_group_members(pgid), + "SIGTERM-immune child survived the SIGKILL escalation", + ) finally: proc.stdout.close() try: From 3d2640cabd7bce85bd03f737644c4866f6678645 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Wed, 8 Jul 2026 09:41:01 -0700 Subject: [PATCH 10/14] [SPARK-57787][CONNECT] Use Markdown code fences in the reuse docs section Review feedback: prefer fenced code blocks over Liquid highlight tags in the new documentation section; they render the same and keep raw-Markdown editors from mistaking comment hashes for headings. Co-authored-by: Isaac --- docs/spark-connect-overview.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/spark-connect-overview.md b/docs/spark-connect-overview.md index be0d2ee31376b..10df9ba448207 100644 --- a/docs/spark-connect-overview.md +++ b/docs/spark-connect-overview.md @@ -281,10 +281,10 @@ The connection may also be programmatically created using _SparkSession#builder_ When you develop or test locally with -{% highlight python %} +```python from pyspark.sql import SparkSession spark = SparkSession.builder.remote("local[*]").getOrCreate() -{% endhighlight %} +``` PySpark boots a fresh in-process Spark Connect server in **every** process. Each `python script.py` run (or each forked test JVM) therefore re-pays the one-time startup cost -- @@ -297,7 +297,7 @@ There are two ways to amortize that cost across runs by reconnecting to a long-l Start one persistent local Spark Connect server and point every run at it: -{% highlight bash %} +```bash # Start once; it stays up across runs. $SPARK_HOME/sbin/start-connect-server.sh --master "local[*]" @@ -306,7 +306,7 @@ python -c 'from pyspark.sql import SparkSession; SparkSession.builder.remote("sc # Stop it when you are done. $SPARK_HOME/sbin/stop-connect-server.sh -{% endhighlight %} +``` ### Let PySpark manage the server (opt-in) @@ -315,11 +315,11 @@ manage a server by hand, enable the opt-in reuse path. The first run starts a ** Connect server and records it in a discovery file; later runs reconnect to it in a fraction of a second: -{% highlight bash %} +```bash export SPARK_LOCAL_CONNECT_REUSE=1 # or .config("spark.local.connect.reuse", "true") python script.py # 1st run: starts a persistent server (cold start, once) python script.py # 2nd+ run: reconnects to it (sub-second) -{% endhighlight %} +``` This is **off by default**; nothing changes unless you opt in. A few details: From 4fa2e114871ee5be0f41a8d78865fbe7356c6de2 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Wed, 8 Jul 2026 15:49:13 -0700 Subject: [PATCH 11/14] [SPARK-57787][CONNECT] Move the reuse machinery into local_server and run the daemon via -m Addresses review comments: - Launch the daemon with `python -m pyspark.sql.connect.local_server` instead of by file path, dropping the sys.path workaround; the daemon always runs from the client's environment, which has the Connect client dependencies. - Move all client-side reuse helpers off SparkSession (they were all staticmethods) into pyspark/sql/connect/local_server.py as module functions; connect/session.py is back to its upstream state. - Pass the auth token only through SPARK_CONNECT_AUTHENTICATE_TOKEN (same precedence as the in-process path); it no longer appears on the daemon argv (visible in ps) nor as a duplicate conf (the server falls back to the env var when the conf is unset). - Probe pid liveness only on POSIX: os.kill(pid, 0) terminates the target process on Windows. - Close the remove-vs-publish race on the discovery file by claiming the start-up lock non-blockingly in _remove_discovery_if_ours. - Turn the start-up lock into a contextmanager; validate discovery value types on read and index the dict directly afterwards; simplify _discovery_path and the makedirs guards; drop the SIGINT try/except (main only runs on the main thread); reap the SIGKILLed daemon with a short wait instead of a bare poll. - Lower the default idle timeout to 1800s, add --stop to the module CLI, and update the docs accordingly. Co-authored-by: Isaac --- docs/spark-connect-overview.md | 8 +- python/pyspark/sql/connect/local_server.py | 526 ++++++++++++++++-- python/pyspark/sql/connect/session.py | 368 ------------ python/pyspark/sql/session.py | 10 +- .../connect/test_connect_local_server.py | 175 +++--- 5 files changed, 586 insertions(+), 501 deletions(-) diff --git a/docs/spark-connect-overview.md b/docs/spark-connect-overview.md index 10df9ba448207..879a59aac9901 100644 --- a/docs/spark-connect-overview.md +++ b/docs/spark-connect-overview.md @@ -336,10 +336,10 @@ This is **off by default**; nothing changes unless you opt in. A few details: concurrent startups can race; callers that lose the race reconnect to the server recorded in the discovery file. - The server self-terminates once it has been idle for `spark.local.connect.server.idleTimeout` - seconds (default `3600`; set `0` to disable). To stop it sooner, call - `SparkSession._stop_local_connect_server()` or terminate the `pid` recorded in the discovery file. - If an old discovery file is rejected after a Spark upgrade, stop the old recorded `pid` if you do - not want to wait for the idle timeout. + seconds (default `1800`, i.e. 30 minutes; set `0` to disable). To stop it sooner, run + `python -m pyspark.sql.connect.local_server --stop` or terminate the `pid` recorded in the + discovery file. If an old discovery file is rejected after a Spark upgrade, stop the old recorded + `pid` if you do not want to wait for the idle timeout. ## Use Spark Connect in standalone applications diff --git a/python/pyspark/sql/connect/local_server.py b/python/pyspark/sql/connect/local_server.py index 19a9b74830d9a..f1d07114f745f 100644 --- a/python/pyspark/sql/connect/local_server.py +++ b/python/pyspark/sql/connect/local_server.py @@ -16,46 +16,82 @@ # """ -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. -""" +Opt-in reuse of a persistent local Spark Connect server (``spark.local.connect.reuse`` / +``SPARK_LOCAL_CONNECT_REUSE``). -import sys +By default ``SparkSession.builder.remote("local[*]").getOrCreate()`` boots a fresh in-process +Connect server in every Python process (see ``SparkSession._start_connect_server`` in +``pyspark.sql.connect.session``). When reuse is enabled, ``reuse_or_start_local_connect_server`` +launches this module once as a detached daemon (``python -m pyspark.sql.connect.local_server``). +The daemon starts a regular (classic) local Spark session with the Spark Connect plugin -- the +same mechanism the in-process path uses -- and then blocks, so one warm JVM and Connect server +keep 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 accepts connections, the daemon writes a *discovery file* recording the host, the +actually bound port, the auth token, its pid and the Spark version. Later client processes read +it and -- after validating that the version matches, the pid is alive and the port accepts +connections -- reconnect instead of starting another server. + +To stop a running server:: -# 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] + python -m pyspark.sql.connect.local_server --stop +""" import argparse +import contextlib import json import os import signal +import socket +import subprocess +import sys +import tempfile import time -from typing import Any +import uuid +from typing import Any, Dict, Iterator, Optional + +from pyspark.errors import PySparkRuntimeError + +# -- the discovery file, shared between the client-side helpers and the daemon -------------------- + + +def _discovery_path() -> str: + """Location of the discovery file describing the running persistent local server.""" + return os.environ.get( + "SPARK_LOCAL_CONNECT_DISCOVERY", + os.path.join(os.path.expanduser("~"), ".spark", "connect-local.json"), + ) + + +def _read_discovery() -> Optional[Dict[str, Any]]: + """Read and validate the discovery file, returning ``None`` if it is absent or malformed. + + A returned dict always has string ``host``, ``token`` and ``spark_version`` values and int + ``port`` and ``pid`` values, so callers can index it without re-validating. + """ + try: + with open(_discovery_path(), "r") as f: + disc = json.load(f) + except (OSError, ValueError): + return None + if not isinstance(disc, dict): + return None + try: + disc["port"] = int(disc["port"]) + disc["pid"] = int(disc["pid"]) + except (KeyError, TypeError, ValueError): + return None + if not all(isinstance(disc.get(k), str) for k in ("host", "token", "spark_version")): + return None + return disc 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): + if parent: os.makedirs(parent, exist_ok=True) payload = { "host": host, @@ -72,17 +108,378 @@ def _write_discovery(path: str, host: str, port: int, token: str, version: str) def _remove_discovery_if_ours(path: str) -> None: - """Remove the discovery file, but only if it still points at this process.""" + """Remove the discovery file at daemon shutdown, but only if it still points at this process. + + The read-check-remove sequence could race with a new server publishing itself and delete the + newcomer's file. Every discovery write happens while a launching client holds the start-up + lock (the daemon writes the file during ``_launch_server``'s lock-held wait), so claiming that + lock non-blockingly closes the race: if it is busy, a new server is starting and will + overwrite the file anyway, and skipping the removal keeps its entry intact. Where ``fcntl`` + is unavailable the removal stays best-effort and unguarded. + """ + lock_fd = None try: - with open(path, "r") as f: - disc = json.load(f) - except (OSError, ValueError): - return - if disc.get("pid") == os.getpid(): + import fcntl + + lock_fd = os.open(path + ".lock", os.O_RDWR | os.O_CREAT, 0o600) + fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except ImportError: + pass # no fcntl (e.g. Windows): fall through to the unguarded removal + except OSError: + if lock_fd is not None: + # The lock is held: a launcher is publishing a new server right now, and the file + # (or what it is about to become) is no longer ours to remove. + os.close(lock_fd) + return + # The lock file itself could not be created: fall through to the unguarded removal. + try: + 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 + finally: + if lock_fd is not None: + os.close(lock_fd) # closing the fd releases the flock + + +# -- client side: decide between reusing the recorded server and starting a fresh one ------------- + + +def _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, that the recorded + process is still alive, and that it is accepting connections on the recorded port. A version + mismatch, dead pid, or closed port means the caller must start its own server instead. The + pid probe runs only on POSIX: on Windows ``os.kill(pid, 0)`` *terminates* the target process + rather than testing it, so there the port probe is the only liveness signal. + """ + from pyspark.version import __version__ + + if disc["spark_version"] != __version__: + return False + if os.name == "posix": try: - os.remove(path) + os.kill(disc["pid"], 0) + except ProcessLookupError: + 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"], disc["port"])) != 0: + return False + return True + + +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 = _read_discovery() + if disc is not None and _server_is_reusable(disc): + os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = disc["token"] + return "sc://{}:{}".format(disc["host"], disc["port"]) + return None + + +@contextlib.contextmanager +def _start_lock() -> Iterator[None]: + """Exclusive cross-process file lock serializing persistent-server start-up. + + On platforms without ``fcntl`` this is a no-op: racing callers may each start a daemon, and + the losers reconnect to the server that wins the discovery-file update. + """ + try: + import fcntl + except ImportError: + yield + return + path = _discovery_path() + ".lock" + parent = os.path.dirname(path) + if parent: + os.makedirs(parent, exist_ok=True) + fd = os.open(path, os.O_RDWR | os.O_CREAT, 0o600) + try: + fcntl.flock(fd, fcntl.LOCK_EX) + yield + finally: + os.close(fd) # closing the fd releases the flock + + +def _port_available(port: int) -> bool: + """Whether ``port`` can currently be bound on localhost (best effort, subject to races).""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + try: + sock.bind(("localhost", port)) + return True + except OSError: + return False + + +def _seed_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 ``SparkSession._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 + + +def _signal_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 + + +def _terminate_server(proc: Any) -> None: + """Terminate a daemon started by ``_launch_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 _signal_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_server (whose getpgid lookup needs a live leader). Signalling this group is safe: + # it was created 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 + try: + # Reap the SIGKILLed daemon so it is not left behind as a zombie. + proc.wait(timeout=5) + except Exception: + pass + + +def _launch_server(master: str, opts: Dict[str, Any]) -> str: + """Launch a detached persistent local Connect server and wait until it is reachable. + + Callers must hold ``_start_lock`` (see ``reuse_or_start_local_connect_server``). If the + server cannot be started but another process has meanwhile published a reusable one, this + reconnects to it rather than failing. + """ + from pyspark.sql.connect.client import DefaultChannelBuilder + + discovery_path = _discovery_path() + # Same token precedence as the in-process ``_start_connect_server``: an explicit env token, + # then one passed as a conf, then a fresh one. + token = ( + os.environ.get("SPARK_CONNECT_AUTHENTICATE_TOKEN") + or opts.get("spark.connect.authenticate.token") + or str(uuid.uuid4()) + ) + + # Choose the port. Tests use an ephemeral port (0) so they can run in parallel. Otherwise we + # honor the configured/default port, but fall back to an ephemeral one if it is already + # taken -- e.g. by a stale server we just rejected on version mismatch -- so a fresh server + # can still start instead of failing to bind. + if "SPARK_TESTING" in os.environ: + port = 0 + else: + port = int( + opts.get("spark.local.connect.server.port", DefaultChannelBuilder.default_port()) + ) + if port != 0 and not _port_available(port): + port = 0 + idle_timeout = opts.get("spark.local.connect.server.idleTimeout", "1800") + + # Seed the server with the caller's start-up confs (warehouse dir, jars, catalog, etc.) so + # first-run behavior matches the in-process path. Passed as a JSON file since confs are + # arbitrary key/values. Written under the discovery dir with 0600 perms as it may hold + # sensitive values, and removed once the daemon has started. + conf_file = None + seed_conf = _seed_conf(opts) + if seed_conf: + parent = os.path.dirname(discovery_path) + if parent: + os.makedirs(parent, exist_ok=True) + fd, conf_file = tempfile.mkstemp(prefix="connect-local-conf-", dir=parent or None) + with os.fdopen(fd, "w") as f: + json.dump(seed_conf, f) + os.chmod(conf_file, 0o600) + + cmd = [ + sys.executable, + "-m", + "pyspark.sql.connect.local_server", + "--master", + master, + "--port", + str(port), + "--discovery", + discovery_path, + "--idle-timeout", + str(idle_timeout), + ] + if conf_file is not None: + cmd += ["--conf-file", conf_file] + + # Launch detached so the server outlives this client process. The token travels through the + # environment, never argv, where it would be visible in `ps` output. + env = dict(os.environ) + env["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = token + for var in ("SPARK_REMOTE", "SPARK_LOCAL_REMOTE", "SPARK_CONNECT_MODE_ENABLED"): + env.pop(var, None) + popen_kwargs: Dict[str, Any] = dict( + env=env, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + if os.name == "posix": + popen_kwargs["start_new_session"] = True + else: + detached = getattr(subprocess, "DETACHED_PROCESS", 0) + new_group = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0) + popen_kwargs["creationflags"] = detached | new_group + proc = subprocess.Popen(cmd, **popen_kwargs) + + try: + # Wait for the server to write the discovery file (which records the actual bound port) + # and start accepting connections. + deadline = time.time() + 120 + while time.time() < deadline: + exit_code = proc.poll() + if exit_code is not None: + # Our daemon died. Another process may have published a usable server in the + # meantime (e.g. a port race), so prefer reconnecting to it over failing. + endpoint = _reuse_from_discovery() + if endpoint is not None: + return endpoint + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={ + "reason": "the server process exited with code {}".format(exit_code) + }, + ) + disc = _read_discovery() + if disc is not None and disc["pid"] == proc.pid and disc["token"] == token: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(0.5) + if sock.connect_ex((disc["host"], disc["port"])) == 0: + os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = token + return "sc://{}:{}".format(disc["host"], disc["port"]) + time.sleep(0.25) + + # Timed out. The daemon may still be inside getOrCreate() -- i.e. before it has wired up + # its own SIGTERM handler -- and it has already spawned a child JVM. Signal the whole + # process group (the daemon is a session leader via start_new_session) and escalate to + # SIGKILL, so the JVM is reaped rather than orphaned. + _terminate_server(proc) + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={"reason": "the server did not become ready within 120 seconds"}, + ) + finally: + if conf_file is not None: + try: + os.remove(conf_file) + except OSError: + pass + + +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 + ``SparkSession._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 = _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. + with _start_lock(): + endpoint = _reuse_from_discovery() + if endpoint is not None: + return endpoint + return _launch_server(master, opts) + + +def stop_local_connect_server() -> bool: + """Stop the persistent local Spark Connect server started by the reuse path, if any. + + Returns ``True`` if a running server was signalled to stop. Safe to call when none is + running. Also exposed on the command line for the dev loop:: + + python -m pyspark.sql.connect.local_server --stop + """ + disc = _read_discovery() + stopped = False + if disc is not None and disc["pid"] != os.getpid(): + stopped = _signal_server(disc["pid"], signal.SIGTERM) + try: + os.remove(_discovery_path()) + except OSError: + pass + return stopped + + +# -- the daemon itself, run as `python -m pyspark.sql.connect.local_server` ----------------------- def _has_active_sessions(spark: Any) -> bool: @@ -115,11 +512,18 @@ def _bound_port(spark: Any, requested_port: int) -> int: def main() -> None: - parser = argparse.ArgumentParser() + parser = argparse.ArgumentParser(description="Persistent local Spark Connect server.") + parser.add_argument( + "--stop", action="store_true", help="stop the recorded running server, if any, and exit" + ) 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( + "--discovery", + default=None, + help="path of the discovery file; defaults to $SPARK_LOCAL_CONNECT_DISCOVERY " + "or ~/.spark/connect-local.json", + ) parser.add_argument( "--conf-file", default=None, @@ -128,24 +532,43 @@ def main() -> None: parser.add_argument( "--idle-timeout", type=float, - default=3600.0, + default=1800.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() + if args.discovery: + # Everything below resolves the path via _discovery_path (as do the client-side helpers, + # e.g. under --stop), so map an explicit --discovery onto the override it honors. + os.environ["SPARK_LOCAL_CONNECT_DISCOVERY"] = args.discovery + discovery = _discovery_path() + + if args.stop: + if stop_local_connect_server(): + print("Stopped the persistent local Spark Connect server.") + else: + print("No running persistent local Spark Connect server found.") + return + # 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 + + # The launcher passes the auth token through the environment (argv would leak it to `ps`); + # generate one for manual runs. The child JVM inherits the env var, and the Connect server + # falls back to it when the `spark.connect.authenticate.token` conf is not set. + token = os.environ.get("SPARK_CONNECT_AUTHENTICATE_TOKEN") + if not token: + token = str(uuid.uuid4()) + os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = 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. + # 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(): @@ -158,28 +581,23 @@ def main() -> None: .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__) + _write_discovery(discovery, "localhost", bound_port, token, __version__) print( "SPARK-CONNECT-LOCAL-SERVER READY port={} pid={}".format(bound_port, os.getpid()), flush=True, ) - stop = {"flag": False} + stop = False def _handle(_signum: int, _frame: Any) -> None: - stop["flag"] = True + nonlocal stop + stop = 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. - pass + signal.signal(signal.SIGINT, _handle) idle_timeout = args.idle_timeout poll_interval = max(1.0, args.poll_interval) @@ -187,7 +605,7 @@ def _handle(_signum: int, _frame: Any) -> None: last_poll = 0.0 try: # Sleep in short slices so a SIGTERM is observed promptly regardless of the poll interval. - while not stop["flag"]: + while not stop: time.sleep(1.0) if idle_timeout <= 0: continue @@ -204,7 +622,7 @@ def _handle(_signum: int, _frame: Any) -> None: elif now - last_active > idle_timeout: break finally: - _remove_discovery_if_ours(args.discovery) + _remove_discovery_if_ours(discovery) spark.stop() diff --git a/python/pyspark/sql/connect/session.py b/python/pyspark/sql/connect/session.py index 126e1ed5652bd..99da7308a1a00 100644 --- a/python/pyspark/sql/connect/session.py +++ b/python/pyspark/sql/connect/session.py @@ -17,11 +17,9 @@ import uuid import json -import signal import threading import os import sys -import time import warnings from collections.abc import Callable, Sized import functools @@ -1249,372 +1247,6 @@ 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() - - @staticmethod - def _start_persistent_local_connect_server(master: str, opts: Dict[str, Any]) -> str: - """Launch a detached persistent local Connect server and wait until it is reachable. - - Callers must hold the start-up lock (see ``_reuse_or_start_local_connect_server``). If the - server cannot be started but another process has meanwhile published a reusable one, this - reconnects to it rather than failing. - """ - import socket - import subprocess - import tempfile - - discovery_path = SparkSession._local_connect_discovery_path() - token = opts.get("spark.connect.authenticate.token") or str(uuid.uuid4()) - - # Choose the port. Tests use an ephemeral port (0) so they can run in parallel. Otherwise we - # honor the configured/default port, but fall back to an ephemeral one if it is already - # taken -- e.g. by a stale server we just rejected on version mismatch -- so a fresh server - # can still start instead of failing to bind. - if "SPARK_TESTING" in os.environ: - port = 0 - else: - port = int( - opts.get("spark.local.connect.server.port", DefaultChannelBuilder.default_port()) - ) - if port != 0 and not SparkSession._local_port_available(port): - port = 0 - idle_timeout = opts.get("spark.local.connect.server.idleTimeout", "3600") - - # Seed the server with the caller's start-up confs (warehouse dir, jars, catalog, etc.) so - # first-run behavior matches the in-process path. Passed as a JSON file since confs are - # arbitrary key/values. Written under the discovery dir with 0600 perms as it may hold - # sensitive values, and removed once the daemon has started. - conf_file = None - seed_conf = SparkSession._local_connect_server_conf(opts) - if seed_conf: - parent = os.path.dirname(discovery_path) - if parent and not os.path.isdir(parent): - os.makedirs(parent, exist_ok=True) - fd, conf_file = tempfile.mkstemp(prefix="connect-local-conf-", dir=parent or None) - with os.fdopen(fd, "w") as f: - json.dump(seed_conf, f) - os.chmod(conf_file, 0o600) - - daemon = os.path.join(os.path.dirname(os.path.abspath(__file__)), "local_server.py") - cmd = [ - sys.executable, - daemon, - "--master", - master, - "--port", - str(port), - "--token", - token, - "--discovery", - discovery_path, - "--idle-timeout", - str(idle_timeout), - ] - if conf_file is not None: - cmd += ["--conf-file", conf_file] - - # Launch detached so the server outlives this client process. - env = dict(os.environ) - for var in ("SPARK_REMOTE", "SPARK_LOCAL_REMOTE", "SPARK_CONNECT_MODE_ENABLED"): - env.pop(var, None) - popen_kwargs: Dict[str, Any] = dict( - env=env, - stdin=subprocess.DEVNULL, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - if os.name == "posix": - popen_kwargs["start_new_session"] = True - else: - detached = getattr(subprocess, "DETACHED_PROCESS", 0) - new_group = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0) - popen_kwargs["creationflags"] = detached | new_group - proc = subprocess.Popen(cmd, **popen_kwargs) - - try: - # Wait for the server to write the discovery file (which records the actual bound port) - # and start accepting connections. - deadline = time.time() + 120 - while time.time() < deadline: - exit_code = proc.poll() - if exit_code is not None: - # Our daemon died. Another process may have published a usable server in the - # meantime (e.g. a port race), so prefer reconnecting to it over failing. - endpoint = SparkSession._reuse_from_discovery() - if endpoint is not None: - return endpoint - raise PySparkRuntimeError( - errorClass="LOCAL_CONNECT_SERVER_START_FAILED", - messageParameters={ - "reason": "the server process exited with code {}".format(exit_code) - }, - ) - disc = SparkSession._read_local_connect_discovery() - if disc is not None and disc.get("pid") == proc.pid and disc.get("token") == token: - 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: - os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = token - return "sc://{}:{}".format(disc["host"], disc["port"]) - time.sleep(0.25) - - # Timed out. The daemon may still be inside getOrCreate() -- i.e. before it has wired up - # its own SIGTERM handler -- and it has already spawned a child JVM. Signal the whole - # process group (the daemon is a session leader via start_new_session) and escalate to - # SIGKILL, so the JVM is reaped rather than orphaned. - SparkSession._terminate_local_connect_server(proc) - raise PySparkRuntimeError( - errorClass="LOCAL_CONNECT_SERVER_START_FAILED", - messageParameters={"reason": "the server did not become ready within 120 seconds"}, - ) - finally: - if conf_file is not None: - try: - os.remove(conf_file) - except OSError: - pass - - @staticmethod - def _stop_local_connect_server() -> bool: - """Stop the persistent local Connect server started by the reuse path, if any. - - Returns ``True`` if a running server was signalled to stop. Safe to call when none is - running. Intended for the dev/test loop, e.g.:: - - python -c "from pyspark.sql.connect.session import SparkSession; \\ - SparkSession._stop_local_connect_server()" - """ - disc = SparkSession._read_local_connect_discovery() - path = SparkSession._local_connect_discovery_path() - stopped = False - if disc is not None: - try: - pid = int(disc["pid"]) - except (TypeError, ValueError, KeyError): - pid = None - if pid is not None and pid != os.getpid(): - stopped = SparkSession._signal_local_connect_server(pid, signal.SIGTERM) - try: - os.remove(path) - except OSError: - pass - return stopped - @property def session_id(self) -> str: return self._session_id diff --git a/python/pyspark/sql/session.py b/python/pyspark/sql/session.py index 32f8f7659b66a..2466e75da693d 100644 --- a/python/pyspark/sql/session.py +++ b/python/pyspark/sql/session.py @@ -531,12 +531,14 @@ def getOrCreate(self) -> "SparkSession": ).lower() in ("1", "true") if url.startswith("local") and reuse_local: + from pyspark.sql.connect.local_server import ( + reuse_or_start_local_connect_server, + ) + # Opt-in: reconnect to a persistent local Connect server (starting # one on the first run) instead of booting a fresh in-process server - # every process. See `_reuse_or_start_local_connect_server`. - url = RemoteSparkSession._reuse_or_start_local_connect_server( - url, opts - ) + # every process. See `pyspark.sql.connect.local_server`. + url = reuse_or_start_local_connect_server(url, opts) for k in ( "spark.local.connect.reuse", "spark.local.connect.server.port", diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server.py b/python/pyspark/sql/tests/connect/test_connect_local_server.py index a43a0388df254..7793add5ad945 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server.py @@ -32,6 +32,7 @@ if should_test_connect: from pyspark.sql import SparkSession as PySparkSession + from pyspark.sql.connect import local_server from pyspark.sql.connect.session import SparkSession as RemoteSparkSession from pyspark.version import __version__ @@ -58,11 +59,11 @@ def tearDown(self) -> None: try: # Only stop a real, separately-spawned server. The discovery-logic unit tests fabricate # discovery files that point at this very process, which must never be signalled. - disc = RemoteSparkSession._read_local_connect_discovery() - if disc is not None and disc.get("pid") != os.getpid(): - port = int(disc["port"]) - RemoteSparkSession._stop_local_connect_server() - # _stop_local_connect_server only signals the daemon and returns; wait for the JVM + disc = local_server._read_discovery() + if disc is not None and disc["pid"] != os.getpid(): + port = disc["port"] + local_server.stop_local_connect_server() + # stop_local_connect_server only signals the daemon and returns; wait for the JVM # to actually release the port so the next test starts from a clean slate. self._wait_port_closed(disc["host"], port) finally: @@ -78,22 +79,22 @@ def tearDown(self) -> None: # -- discovery / reuse-decision logic (no real server) ---------------------------------------- def test_discovery_path_honors_override(self) -> None: - self.assertEqual(RemoteSparkSession._local_connect_discovery_path(), self._discovery) + self.assertEqual(local_server._discovery_path(), self._discovery) os.environ.pop("SPARK_LOCAL_CONNECT_DISCOVERY") self.assertTrue( - RemoteSparkSession._local_connect_discovery_path().endswith( - os.path.join(".spark", "connect-local.json") - ) + local_server._discovery_path().endswith(os.path.join(".spark", "connect-local.json")) ) def test_read_discovery_missing_or_malformed(self) -> None: - self.assertIsNone(RemoteSparkSession._read_local_connect_discovery()) + self.assertIsNone(local_server._read_discovery()) with open(self._discovery, "w") as f: f.write("not json") - self.assertIsNone(RemoteSparkSession._read_local_connect_discovery()) + self.assertIsNone(local_server._read_discovery()) with open(self._discovery, "w") as f: json.dump({"host": "localhost"}, f) # missing required keys - self.assertIsNone(RemoteSparkSession._read_local_connect_discovery()) + self.assertIsNone(local_server._read_discovery()) + self._write_discovery(pid="not-a-pid") # all keys present, but pid is not an int + self.assertIsNone(local_server._read_discovery()) def _write_discovery(self, **overrides) -> dict: disc = { @@ -110,12 +111,12 @@ def _write_discovery(self, **overrides) -> dict: def test_not_reusable_on_version_mismatch(self) -> None: disc = self._write_discovery(spark_version="0.0.0-not-this-build") - self.assertFalse(RemoteSparkSession._local_connect_server_is_reusable(disc)) + self.assertFalse(local_server._server_is_reusable(disc)) def test_not_reusable_on_dead_pid(self) -> None: # PID 2**31 - 1 is effectively guaranteed not to exist. disc = self._write_discovery(pid=2**31 - 1, port=1) - self.assertFalse(RemoteSparkSession._local_connect_server_is_reusable(disc)) + self.assertFalse(local_server._server_is_reusable(disc)) def test_reusable_when_alive_and_listening(self) -> None: # A live listening socket owned by this (alive) process with a matching version is reusable. @@ -125,31 +126,60 @@ def test_reusable_when_alive_and_listening(self) -> None: listener.listen(1) port = listener.getsockname()[1] disc = self._write_discovery(port=port) - self.assertTrue(RemoteSparkSession._local_connect_server_is_reusable(disc)) + self.assertTrue(local_server._server_is_reusable(disc)) finally: listener.close() # Once the socket is closed the port is no longer reachable, so it is not reusable. - self.assertFalse(RemoteSparkSession._local_connect_server_is_reusable(disc)) + self.assertFalse(local_server._server_is_reusable(disc)) + + def test_pid_probe_is_skipped_on_windows(self) -> None: + """The pid liveness probe must never run on Windows. + + There ``os.kill(pid, 0)`` does not probe the process -- it unconditionally *terminates* + it via TerminateProcess -- so the reuse check would kill the very server it is examining. + """ + from unittest import mock + + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + listener.bind(("localhost", 0)) + listener.listen(1) + disc = self._write_discovery(port=listener.getsockname()[1]) + with mock.patch.object(os, "name", "nt"), mock.patch.object(os, "kill") as kill: + self.assertTrue(local_server._server_is_reusable(disc)) + kill.assert_not_called() + finally: + listener.close() def test_stop_when_no_server_is_safe(self) -> None: - self.assertFalse(RemoteSparkSession._stop_local_connect_server()) + self.assertFalse(local_server.stop_local_connect_server()) def test_stop_signals_recorded_server_group(self) -> None: + from unittest import mock + self._write_discovery(pid=12345) calls = [] - old_signal = RemoteSparkSession._signal_local_connect_server def fake_signal(pid, sig): calls.append((pid, sig)) return True - try: - RemoteSparkSession._signal_local_connect_server = staticmethod(fake_signal) - self.assertTrue(RemoteSparkSession._stop_local_connect_server()) - self.assertEqual(calls, [(12345, signal.SIGTERM)]) - self.assertFalse(os.path.exists(self._discovery)) - finally: - RemoteSparkSession._signal_local_connect_server = old_signal + with mock.patch.object(local_server, "_signal_server", fake_signal): + self.assertTrue(local_server.stop_local_connect_server()) + self.assertEqual(calls, [(12345, signal.SIGTERM)]) + self.assertFalse(os.path.exists(self._discovery)) + + def test_stop_cli_reports_when_no_server(self) -> None: + """`python -m pyspark.sql.connect.local_server --stop` is safe with nothing running.""" + result = subprocess.run( + [sys.executable, "-m", "pyspark.sql.connect.local_server", "--stop"], + env=dict(os.environ), + capture_output=True, + text=True, + timeout=120, + ) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertIn("No running persistent local Spark Connect server", result.stdout) @unittest.skipUnless(os.name == "posix", "process groups are POSIX-only") def test_signal_group_kills_only_session_leaders(self) -> None: @@ -168,9 +198,7 @@ def test_signal_group_kills_only_session_leaders(self) -> None: try: self.assertNotEqual(os.getpgid(child.pid), child.pid) with mock.patch("os.killpg") as killpg, mock.patch("os.kill") as kill: - self.assertTrue( - RemoteSparkSession._signal_local_connect_server(child.pid, signal.SIGTERM) - ) + self.assertTrue(local_server._signal_server(child.pid, signal.SIGTERM)) killpg.assert_not_called() kill.assert_called_once_with(child.pid, signal.SIGTERM) finally: @@ -181,9 +209,7 @@ def test_signal_group_kills_only_session_leaders(self) -> None: try: self.assertEqual(os.getpgid(leader.pid), leader.pid) with mock.patch("os.killpg") as killpg, mock.patch("os.kill") as kill: - self.assertTrue( - RemoteSparkSession._signal_local_connect_server(leader.pid, signal.SIGTERM) - ) + self.assertTrue(local_server._signal_server(leader.pid, signal.SIGTERM)) killpg.assert_called_once_with(leader.pid, signal.SIGTERM) kill.assert_not_called() finally: @@ -191,7 +217,7 @@ def test_signal_group_kills_only_session_leaders(self) -> None: leader.communicate() def test_reuse_from_discovery_none_when_absent(self) -> None: - self.assertIsNone(RemoteSparkSession._reuse_from_discovery()) + self.assertIsNone(local_server._reuse_from_discovery()) def test_local_port_available(self) -> None: listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) @@ -199,14 +225,14 @@ def test_local_port_available(self) -> None: listener.bind(("localhost", 0)) listener.listen(1) taken = listener.getsockname()[1] - self.assertFalse(RemoteSparkSession._local_port_available(taken)) + self.assertFalse(local_server._port_available(taken)) finally: listener.close() # The port is free again once the listener is closed. - self.assertTrue(RemoteSparkSession._local_port_available(taken)) + self.assertTrue(local_server._port_available(taken)) def test_server_conf_seeds_user_confs_and_drops_control_keys(self) -> None: - """_local_connect_server_conf keeps user startup confs but not keys the daemon controls.""" + """_seed_conf keeps user startup confs but not keys the daemon controls.""" opts = { "spark.sql.warehouse.dir": "/tmp/wh", "spark.jars.packages": "org.example:lib:1.0", @@ -218,7 +244,7 @@ def test_server_conf_seeds_user_confs_and_drops_control_keys(self) -> None: "spark.local.connect.server.port": "15002", "spark.local.connect.server.idleTimeout": "60", } - conf = RemoteSparkSession._local_connect_server_conf(opts) + conf = local_server._seed_conf(opts) self.assertEqual(conf.get("spark.sql.warehouse.dir"), "/tmp/wh") self.assertEqual(conf.get("spark.jars.packages"), "org.example:lib:1.0") for dropped in ( @@ -233,13 +259,28 @@ def test_server_conf_seeds_user_confs_and_drops_control_keys(self) -> None: self.assertNotIn(dropped, conf) def test_start_lock_roundtrip(self) -> None: - """Acquiring and releasing the start-up lock creates the lock file and does not error.""" - fd = RemoteSparkSession._acquire_local_connect_start_lock() - try: - if fd is not None: # None only where fcntl is unavailable (e.g. Windows) - self.assertTrue(os.path.exists(self._discovery + ".lock")) - finally: - RemoteSparkSession._release_local_connect_start_lock(fd) + """Entering and exiting the start-up lock creates the lock file and does not error.""" + with local_server._start_lock(): + try: + import fcntl # noqa: F401 + except ImportError: + # Locking is a no-op without fcntl; entering the context is all we can assert. + return + self.assertTrue(os.path.exists(self._discovery + ".lock")) + + @unittest.skipUnless(os.name == "posix", "the removal guard needs fcntl") + def test_remove_discovery_skipped_while_start_lock_held(self) -> None: + """Daemon shutdown must not delete the discovery file while a launcher holds the lock. + + A held start-up lock means another process is publishing a new server; removing the file + in that window could delete the newcomer's entry (see _remove_discovery_if_ours). + """ + self._write_discovery(pid=os.getpid()) + with local_server._start_lock(): + local_server._remove_discovery_if_ours(self._discovery) + self.assertTrue(os.path.exists(self._discovery)) + local_server._remove_discovery_if_ours(self._discovery) + self.assertFalse(os.path.exists(self._discovery)) # -- end-to-end: start a real detached server, reconnect to it, verify isolation -------------- @@ -275,7 +316,7 @@ def test_builder_remote_local_uses_reuse_flag(self) -> None: ) self.assertEqual(spark.range(2).count(), 2) - disc = RemoteSparkSession._read_local_connect_discovery() + disc = local_server._read_discovery() self.assertIsNotNone(disc) self.assertEqual(disc["spark_version"], __version__) self.assertNotEqual(disc["pid"], os.getpid()) @@ -284,8 +325,7 @@ def test_builder_remote_local_uses_reuse_flag(self) -> None: spark.stop() def test_concurrent_startup_reuses_one_server(self) -> None: - script = textwrap.dedent( - """ + script = textwrap.dedent(""" import json import os @@ -303,8 +343,7 @@ def test_concurrent_startup_reuses_one_server(self) -> None: print(json.dumps({"count": count, "pid": disc["pid"], "port": disc["port"]})) finally: spark.stop() - """ - ) + """) env = dict(os.environ) env["SPARK_LOCAL_CONNECT_DISCOVERY"] = self._discovery env["SPARK_LOCAL_CONNECT_REUSE"] = "1" @@ -339,10 +378,10 @@ def test_concurrent_startup_reuses_one_server(self) -> None: def test_start_reuse_and_session_isolation(self) -> None: # First call starts a detached persistent server and records it in the discovery file. - endpoint = RemoteSparkSession._reuse_or_start_local_connect_server("local[2]", {}) + endpoint = local_server.reuse_or_start_local_connect_server("local[2]", {}) self.assertTrue(endpoint.startswith("sc://localhost:")) - disc = RemoteSparkSession._read_local_connect_discovery() + disc = local_server._read_discovery() self.assertIsNotNone(disc) self.assertEqual(disc["spark_version"], __version__) self.assertEqual(os.environ.get("SPARK_CONNECT_AUTHENTICATE_TOKEN"), disc["token"]) @@ -351,9 +390,9 @@ def test_start_reuse_and_session_isolation(self) -> None: s1 = s2 = None try: # A second call reuses the running server: same endpoint, no new process spawned. - endpoint2 = RemoteSparkSession._reuse_or_start_local_connect_server("local[2]", {}) + endpoint2 = local_server.reuse_or_start_local_connect_server("local[2]", {}) self.assertEqual(endpoint2, endpoint) - self.assertEqual(RemoteSparkSession._read_local_connect_discovery()["pid"], first_pid) + self.assertEqual(local_server._read_discovery()["pid"], first_pid) # Two independent client connections to the same server run real queries... s1 = RemoteSparkSession.builder.remote(endpoint).create() @@ -372,8 +411,8 @@ def test_start_reuse_and_session_isolation(self) -> None: self._release(s2) # Stopping signals the server and removes the discovery file. - self.assertTrue(RemoteSparkSession._stop_local_connect_server()) - self.assertIsNone(RemoteSparkSession._read_local_connect_discovery()) + self.assertTrue(local_server.stop_local_connect_server()) + self.assertIsNone(local_server._read_discovery()) # The server should stop accepting connections shortly afterwards. (We check the port rather # than the pid: the detached server is a child of this long-lived test process, so once it # exits it lingers as an unreaped zombie for which os.kill(pid, 0) still succeeds.) @@ -387,7 +426,7 @@ def test_start_seeds_static_conf_on_the_server(self) -> None: """A start-up conf passed by the first caller reaches the daemon's SparkConf.""" warehouse = os.path.join(self._tmpdir, "seeded-wh") opts = {"spark.sql.warehouse.dir": warehouse} - endpoint = RemoteSparkSession._reuse_or_start_local_connect_server("local[2]", opts) + endpoint = local_server.reuse_or_start_local_connect_server("local[2]", opts) spark = None try: spark = RemoteSparkSession.builder.remote(endpoint).create() @@ -433,15 +472,14 @@ def test_terminate_escalates_when_group_outlives_daemon(self) -> None: Reproduces the CI failure mode of test_terminate_reaps_daemon_and_jvm: the daemon exits on SIGTERM promptly while its JVM is still shutting down (or wedged) past the grace - period. _terminate_local_connect_server must not return while the group has live members. + period. _terminate_server must not return while the group has live members. """ if shutil.which("ps") is None: self.skipTest("ps is needed to inspect process group members") # A session leader that exits on SIGTERM after spawning a SIGTERM-immune child into its # process group -- the child stands in for a JVM that outlives the daemon. - leader_prog = textwrap.dedent( - """ + leader_prog = textwrap.dedent(""" import signal import subprocess import sys @@ -460,8 +498,7 @@ def test_terminate_escalates_when_group_outlives_daemon(self) -> None: print("pid:%d" % child.pid, flush=True) signal.signal(signal.SIGTERM, lambda *_: sys.exit(0)) time.sleep(600) - """ - ) + """) proc = subprocess.Popen( [sys.executable, "-c", leader_prog], stdout=subprocess.PIPE, @@ -475,11 +512,11 @@ def test_terminate_escalates_when_group_outlives_daemon(self) -> None: self.assertIn(b"ready", lines) child_pid = int(next(ln for ln in lines if ln.startswith(b"pid:"))[4:]) - RemoteSparkSession._terminate_local_connect_server(proc) + local_server._terminate_server(proc) self.assertTrue( self._wait_process_group_gone(pgid, timeout=10), - "process group survived _terminate_local_connect_server", + "process group survived _terminate_server", ) # The SIGTERM-immune child no longer runs, proving the SIGKILL escalation fired. self.assertNotIn( @@ -496,7 +533,7 @@ def test_terminate_escalates_when_group_outlives_daemon(self) -> None: @unittest.skipUnless(os.name == "posix", "process-group reaping is POSIX-only") def test_terminate_reaps_daemon_and_jvm(self) -> None: - """_terminate_local_connect_server kills the daemon *and* its child JVM. + """_terminate_server kills the daemon *and* its child JVM. The test waits until the daemon has launched the JVM, then checks that the whole process group exits. @@ -504,18 +541,14 @@ def test_terminate_reaps_daemon_and_jvm(self) -> None: if shutil.which("pgrep") is None: self.skipTest("pgrep is needed to detect the child JVM") - import pyspark.sql.connect.session as session_mod - - daemon = os.path.join(os.path.dirname(session_mod.__file__), "local_server.py") cmd = [ sys.executable, - daemon, + "-m", + "pyspark.sql.connect.local_server", "--master", "local[2]", "--port", "0", # ephemeral, so a stray server elsewhere cannot interfere - "--token", - "reap-test", "--discovery", self._discovery, "--idle-timeout", @@ -554,11 +587,11 @@ def child_pids(): time.sleep(0.5) self.assertTrue(kids, "daemon never launched a JVM; cannot exercise the orphan case") - RemoteSparkSession._terminate_local_connect_server(proc) + local_server._terminate_server(proc) self.assertTrue( self._wait_process_group_gone(pgid, timeout=30), - "process group survived _terminate_local_connect_server", + "process group survived _terminate_server", ) finally: try: From 89f6ad86a721a1e0ef46519b4816cdee2baac531 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Wed, 8 Jul 2026 16:47:52 -0700 Subject: [PATCH 12/14] [SPARK-57787][CONNECT] Start the persistent local server via sbin/start-connect-server.sh Reworks the opt-in local Connect server reuse path per review feedback: instead of a detached Python daemon that boots a classic Py4J session with the Connect plugin and babysits it, the launcher now drives the standard `sbin/start-connect-server.sh` (spark-daemon.sh submit SparkConnectServer) and lets spark-daemon.sh own the JVM (pid file, logs, daemonization). - No Python server process any more: PySpark runs the script with the auth token in its env, waits for the port, writes the discovery file (pid taken from the spark-daemon pid file), and returns the endpoint. - Deletes the Py4J usage, signal handling, idle-session polling, process-group management and SIGKILL escalation, and the discovery-file removal race guard -- none of it has an equivalent in the new shape. - Idle self-termination goes away with the daemon; the server runs until stopped (`python -m pyspark.sql.connect.local_server --stop`, or `sbin/stop-connect-server.sh`). `spark.local.connect.server.idleTimeout` is removed. - The server port is picked client-side (configured port when free, else an OS-assigned one), since the standalone script cannot report a bound ephemeral port back. - Seed confs travel via a 0600 spark-submit --properties-file instead of a JSON file parsed by the daemon. Artifact isolation needs no special handling: SparkConnectServer enables it itself. - The reuse path is POSIX-only now (it drives the sbin shell scripts) and says so explicitly; server logs land next to the discovery file (~/.spark/logs by default) for debuggability. - Ships start/stop-connect-server.sh in the pip sbin package (their dependencies, bin/* and spark-daemon.sh, were already packaged). Net effect vs the previous approach: -313 lines, and the server lifecycle is exactly the one users get from the documented manual workflow. Co-authored-by: Isaac --- docs/spark-connect-overview.md | 19 +- python/packaging/classic/setup.py | 2 + python/pyspark/sql/connect/local_server.py | 590 +++++++----------- python/pyspark/sql/session.py | 9 +- .../connect/test_connect_local_server.py | 273 ++------ 5 files changed, 290 insertions(+), 603 deletions(-) diff --git a/docs/spark-connect-overview.md b/docs/spark-connect-overview.md index 879a59aac9901..72896b0bfa83f 100644 --- a/docs/spark-connect-overview.md +++ b/docs/spark-connect-overview.md @@ -311,9 +311,10 @@ $SPARK_HOME/sbin/stop-connect-server.sh ### Let PySpark manage the server (opt-in) If you would rather keep your code as `SparkSession.builder.remote("local[*]").getOrCreate()` and not -manage a server by hand, enable the opt-in reuse path. The first run starts a **detached** local -Connect server and records it in a discovery file; later runs reconnect to it in a fraction of a -second: +manage a server by hand, enable the opt-in reuse path. The first run starts one persistent local +Connect server through the standard `sbin/start-connect-server.sh` script -- the same daemon you +would start by hand above -- and records it in a discovery file; later runs reconnect to it in a +fraction of a second: ```bash export SPARK_LOCAL_CONNECT_REUSE=1 # or .config("spark.local.connect.reuse", "true") @@ -335,11 +336,13 @@ This is **off by default**; nothing changes unless you opt in. A few details: Unix-like systems, PySpark uses a file lock around first startup. On platforms without `fcntl`, concurrent startups can race; callers that lose the race reconnect to the server recorded in the discovery file. -- The server self-terminates once it has been idle for `spark.local.connect.server.idleTimeout` - seconds (default `1800`, i.e. 30 minutes; set `0` to disable). To stop it sooner, run - `python -m pyspark.sql.connect.local_server --stop` or terminate the `pid` recorded in the - discovery file. If an old discovery file is rejected after a Spark upgrade, stop the old recorded - `pid` if you do not want to wait for the idle timeout. +- The server is a regular `spark-daemon.sh`-managed process and runs until you stop it with + `python -m pyspark.sql.connect.local_server --stop` (or `sbin/stop-connect-server.sh`, or by + terminating the `pid` recorded in the discovery file). Its pid file and logs live next to the + discovery file (`~/.spark` and `~/.spark/logs` by default), which is where to look if a start-up + fails. After a Spark upgrade the old server is rejected on its recorded version; stop it the same + way. Because this path drives the `sbin/` shell scripts, it is unavailable on Windows -- there, + start a server manually and connect to it with `.remote("sc://...")`. ## Use Spark Connect in standalone applications diff --git a/python/packaging/classic/setup.py b/python/packaging/classic/setup.py index 53d11a917f55e..1b6df412589b7 100755 --- a/python/packaging/classic/setup.py +++ b/python/packaging/classic/setup.py @@ -345,7 +345,9 @@ def run(self): "pyspark.sbin": [ "spark-config.sh", "spark-daemon.sh", + "start-connect-server.sh", "start-history-server.sh", + "stop-connect-server.sh", "stop-history-server.sh", ], "pyspark.python.lib": ["*.zip"], diff --git a/python/pyspark/sql/connect/local_server.py b/python/pyspark/sql/connect/local_server.py index f1d07114f745f..9019750a06420 100644 --- a/python/pyspark/sql/connect/local_server.py +++ b/python/pyspark/sql/connect/local_server.py @@ -21,22 +21,22 @@ By default ``SparkSession.builder.remote("local[*]").getOrCreate()`` boots a fresh in-process Connect server in every Python process (see ``SparkSession._start_connect_server`` in -``pyspark.sql.connect.session``). When reuse is enabled, ``reuse_or_start_local_connect_server`` -launches this module once as a detached daemon (``python -m pyspark.sql.connect.local_server``). -The daemon starts a regular (classic) local Spark session with the Spark Connect plugin -- the -same mechanism the in-process path uses -- and then blocks, so one warm JVM and Connect server -keep 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 accepts connections, the daemon writes a *discovery file* recording the host, the -actually bound port, the auth token, its pid and the Spark version. Later client processes read -it and -- after validating that the version matches, the pid is alive and the port accepts -connections -- reconnect instead of starting another server. - -To stop a running server:: +``pyspark.sql.connect.session``). When reuse is enabled, the first run instead starts one +long-lived server through the standard ``sbin/start-connect-server.sh`` script -- the same +daemon a user would start by hand -- and records how to reach it in a *discovery file* (host, +port, auth token, pid and Spark version). Later runs validate that record (matching Spark +version, live pid, port accepting connections) and reconnect instead of starting another +server. Each client connection gets its own isolated server-side session, so session-local +state (temp views, runtime SQL confs, session artifacts) does not leak between runs. + +``sbin/spark-daemon.sh`` owns the server process: its pid file and logs are kept next to the +discovery file (``~/.spark`` by default; override with ``SPARK_LOCAL_CONNECT_DISCOVERY``). The +server runs until stopped with:: python -m pyspark.sql.connect.local_server --stop + +or ``sbin/stop-connect-server.sh``. This relies on the POSIX shell scripts under ``sbin/``, so +the reuse path is not supported on Windows. """ import argparse @@ -54,7 +54,13 @@ from pyspark.errors import PySparkRuntimeError -# -- the discovery file, shared between the client-side helpers and the daemon -------------------- +_SERVER_CLASS = "org.apache.spark.sql.connect.service.SparkConnectServer" +# A fixed SPARK_IDENT_STRING keeps the spark-daemon.sh pid and log file names stable +# regardless of $USER. +_SPARK_IDENT = "local-connect" + + +# -- the discovery file -------------------------------------------------------------------------- def _discovery_path() -> str: @@ -65,6 +71,11 @@ def _discovery_path() -> str: ) +def _runtime_dir() -> str: + """Directory holding the discovery file, and with it the server's pid file and logs.""" + return os.path.dirname(_discovery_path()) or os.getcwd() + + def _read_discovery() -> Optional[Dict[str, Any]]: """Read and validate the discovery file, returning ``None`` if it is absent or malformed. @@ -88,7 +99,7 @@ def _read_discovery() -> Optional[Dict[str, Any]]: return disc -def _write_discovery(path: str, host: str, port: int, token: str, version: str) -> None: +def _write_discovery(path: str, host: str, port: int, token: str, pid: int, version: str) -> None: """Atomically write the discovery file with ``0600`` perms (it holds the auth token).""" parent = os.path.dirname(path) if parent: @@ -97,7 +108,7 @@ def _write_discovery(path: str, host: str, port: int, token: str, version: str) "host": host, "port": port, "token": token, - "pid": os.getpid(), + "pid": pid, "spark_version": version, } tmp = "{}.{}.tmp".format(path, os.getpid()) @@ -107,48 +118,18 @@ def _write_discovery(path: str, host: str, port: int, token: str, version: str) os.replace(tmp, path) -def _remove_discovery_if_ours(path: str) -> None: - """Remove the discovery file at daemon shutdown, but only if it still points at this process. +# -- deciding between reusing the recorded server and starting a fresh one ------------------------ - The read-check-remove sequence could race with a new server publishing itself and delete the - newcomer's file. Every discovery write happens while a launching client holds the start-up - lock (the daemon writes the file during ``_launch_server``'s lock-held wait), so claiming that - lock non-blockingly closes the race: if it is busy, a new server is starting and will - overwrite the file anyway, and skipping the removal keeps its entry intact. Where ``fcntl`` - is unavailable the removal stays best-effort and unguarded. - """ - lock_fd = None - try: - import fcntl - lock_fd = os.open(path + ".lock", os.O_RDWR | os.O_CREAT, 0o600) - fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) - except ImportError: - pass # no fcntl (e.g. Windows): fall through to the unguarded removal - except OSError: - if lock_fd is not None: - # The lock is held: a launcher is publishing a new server right now, and the file - # (or what it is about to become) is no longer ours to remove. - os.close(lock_fd) - return - # The lock file itself could not be created: fall through to the unguarded removal. +def _pid_alive(pid: int) -> bool: + """Whether ``pid`` exists (POSIX only). A process we cannot signal counts as alive.""" try: - 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 - finally: - if lock_fd is not None: - os.close(lock_fd) # closing the fd releases the flock - - -# -- client side: decide between reusing the recorded server and starting a fresh one ------------- + os.kill(pid, 0) + except ProcessLookupError: + return False + except OSError: + pass + return True def _server_is_reusable(disc: Dict[str, Any]) -> bool: @@ -164,14 +145,8 @@ def _server_is_reusable(disc: Dict[str, Any]) -> bool: if disc["spark_version"] != __version__: return False - if os.name == "posix": - try: - os.kill(disc["pid"], 0) - except ProcessLookupError: - return False - except OSError: - # The process exists but is not ours to signal (e.g. PermissionError): treat as alive. - pass + if os.name == "posix" and not _pid_alive(disc["pid"]): + return False with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.settimeout(0.5) if sock.connect_ex((disc["host"], disc["port"])) != 0: @@ -196,8 +171,8 @@ def _reuse_from_discovery() -> Optional[str]: def _start_lock() -> Iterator[None]: """Exclusive cross-process file lock serializing persistent-server start-up. - On platforms without ``fcntl`` this is a no-op: racing callers may each start a daemon, and - the losers reconnect to the server that wins the discovery-file update. + On platforms without ``fcntl`` this is a no-op: racing callers may each start a server, and + the losers reconnect to the one that wins the discovery-file update. """ try: import fcntl @@ -216,14 +191,35 @@ def _start_lock() -> Iterator[None]: os.close(fd) # closing the fd releases the flock -def _port_available(port: int) -> bool: - """Whether ``port`` can currently be bound on localhost (best effort, subject to races).""" +# -- starting and stopping the server through sbin/start-connect-server.sh ------------------------ + + +def _pick_port(opts: Dict[str, Any]) -> int: + """Choose the port for a fresh server. + + Tests always use an OS-assigned free port so suites can run in parallel. Otherwise the + configured/default port is honored unless it is already taken -- e.g. by a stale server just + rejected on version mismatch -- in which case an OS-assigned port is used instead. Unlike the + in-process path, the standalone script cannot report an ephemeral port back to us, so the + free port is picked (and released) here, subject to a small race until the server binds it. + """ + + def free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("localhost", 0)) + return sock.getsockname()[1] + + if "SPARK_TESTING" in os.environ: + return free_port() + from pyspark.sql.connect.client import DefaultChannelBuilder + + port = int(opts.get("spark.local.connect.server.port", DefaultChannelBuilder.default_port())) with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: try: sock.bind(("localhost", port)) - return True + return port except OSError: - return False + return free_port() def _seed_conf(opts: Dict[str, Any]) -> Dict[str, Any]: @@ -231,205 +227,183 @@ def _seed_conf(opts: Dict[str, Any]) -> Dict[str, Any]: Mirrors the merge that the in-process ``SparkSession._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. + catalog confs, etc.). Keys the launcher controls itself (master, port, token) and the + ``spark.local.connect.*`` 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) + for k in list(conf): + if k in ( + "spark.remote", + "spark.api.mode", + "spark.master", + "spark.connect.authenticate.token", + "spark.connect.grpc.binding.port", + ) or k.startswith("spark.local.connect."): + conf.pop(k) return conf -def _signal_server(pid: int, sig: int) -> bool: - """Signal a detached local Connect daemon, including its JVM on POSIX. +def _write_seed_properties(opts: Dict[str, Any], runtime_dir: str) -> Optional[str]: + """Write the seed configs as a ``--properties-file`` for spark-submit, or return ``None``. - 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. + Written with ``0600`` perms since configs may hold sensitive values; passing a file keeps + them off the server's argv, where they would be visible in ``ps`` output. """ - 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 + seed = _seed_conf(opts) + if not seed: + return None + fd, path = tempfile.mkstemp(prefix="connect-local-conf-", suffix=".properties", dir=runtime_dir) + with os.fdopen(fd, "w") as f: + for key, value in seed.items(): + escaped = str(value).replace("\\", "\\\\").replace("\n", "\\n") + f.write("{}={}\n".format(key, escaped)) + os.chmod(path, 0o600) + return path -def _terminate_server(proc: Any) -> None: - """Terminate a daemon started by ``_launch_server`` and its JVM. +def _pid_file(runtime_dir: str) -> str: + """The pid file spark-daemon.sh maintains for the server started by this module.""" + return os.path.join(runtime_dir, "spark-{}-{}-1.pid".format(_SPARK_IDENT, _SERVER_CLASS)) - 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 _signal_server(proc.pid, signal.SIGTERM): - return + +def _read_pid(path: str) -> Optional[int]: 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_server (whose getpgid lookup needs a live leader). Signalling this group is safe: - # it was created 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) + with open(path, "r") as f: + return int(f.read().strip()) + except (OSError, ValueError): + return None + + +def _signal_server(pid: int, sig: int) -> bool: + """Best-effort signal to the recorded server pid (the JVM started by spark-daemon.sh).""" try: - os.killpg(proc.pid, signal.SIGKILL) + os.kill(pid, sig) + return True except OSError: - pass - try: - # Reap the SIGKILLed daemon so it is not left behind as a zombie. - proc.wait(timeout=5) - except Exception: - pass + return False def _launch_server(master: str, opts: Dict[str, Any]) -> str: - """Launch a detached persistent local Connect server and wait until it is reachable. + """Start a persistent local Connect server and wait until it is reachable. - Callers must hold ``_start_lock`` (see ``reuse_or_start_local_connect_server``). If the - server cannot be started but another process has meanwhile published a reusable one, this - reconnects to it rather than failing. + Runs the standard ``sbin/start-connect-server.sh``, which daemonizes the server JVM through + ``sbin/spark-daemon.sh`` (pid file, logs). Once the server accepts connections, the + discovery file is written and the ``sc://host:port`` endpoint returned. Callers must hold + ``_start_lock`` (see ``reuse_or_start_local_connect_server``). """ - from pyspark.sql.connect.client import DefaultChannelBuilder + from pyspark.find_spark_home import _find_spark_home + from pyspark.version import __version__ + + spark_home = os.environ.get("SPARK_HOME") or _find_spark_home() + script = os.path.join(spark_home, "sbin", "start-connect-server.sh") + if not os.path.isfile(script): + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={"reason": "cannot find {}".format(script)}, + ) discovery_path = _discovery_path() + runtime_dir = _runtime_dir() + os.makedirs(runtime_dir, exist_ok=True) + log_dir = os.path.join(runtime_dir, "logs") + pid_file = _pid_file(runtime_dir) + # Same token precedence as the in-process ``_start_connect_server``: an explicit env token, - # then one passed as a conf, then a fresh one. + # then one passed as a conf, then a fresh one. The token travels through the environment, + # never argv, where it would be visible in `ps` output. token = ( os.environ.get("SPARK_CONNECT_AUTHENTICATE_TOKEN") or opts.get("spark.connect.authenticate.token") or str(uuid.uuid4()) ) + port = _pick_port(opts) + conf_file = _write_seed_properties(opts, runtime_dir) - # Choose the port. Tests use an ephemeral port (0) so they can run in parallel. Otherwise we - # honor the configured/default port, but fall back to an ephemeral one if it is already - # taken -- e.g. by a stale server we just rejected on version mismatch -- so a fresh server - # can still start instead of failing to bind. - if "SPARK_TESTING" in os.environ: - port = 0 - else: - port = int( - opts.get("spark.local.connect.server.port", DefaultChannelBuilder.default_port()) - ) - if port != 0 and not _port_available(port): - port = 0 - idle_timeout = opts.get("spark.local.connect.server.idleTimeout", "1800") - - # Seed the server with the caller's start-up confs (warehouse dir, jars, catalog, etc.) so - # first-run behavior matches the in-process path. Passed as a JSON file since confs are - # arbitrary key/values. Written under the discovery dir with 0600 perms as it may hold - # sensitive values, and removed once the daemon has started. - conf_file = None - seed_conf = _seed_conf(opts) - if seed_conf: - parent = os.path.dirname(discovery_path) - if parent: - os.makedirs(parent, exist_ok=True) - fd, conf_file = tempfile.mkstemp(prefix="connect-local-conf-", dir=parent or None) - with os.fdopen(fd, "w") as f: - json.dump(seed_conf, f) - os.chmod(conf_file, 0o600) + env = dict(os.environ) + for var in ("SPARK_REMOTE", "SPARK_LOCAL_REMOTE", "SPARK_CONNECT_MODE_ENABLED"): + env.pop(var, None) + env["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = token + env["SPARK_PID_DIR"] = runtime_dir + env["SPARK_LOG_DIR"] = log_dir + env["SPARK_IDENT_STRING"] = _SPARK_IDENT cmd = [ - sys.executable, - "-m", - "pyspark.sql.connect.local_server", + script, "--master", master, - "--port", - str(port), - "--discovery", - discovery_path, - "--idle-timeout", - str(idle_timeout), + "--conf", + "spark.connect.grpc.binding.port={}".format(port), ] if conf_file is not None: - cmd += ["--conf-file", conf_file] - - # Launch detached so the server outlives this client process. The token travels through the - # environment, never argv, where it would be visible in `ps` output. - env = dict(os.environ) - env["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = token - for var in ("SPARK_REMOTE", "SPARK_LOCAL_REMOTE", "SPARK_CONNECT_MODE_ENABLED"): - env.pop(var, None) - popen_kwargs: Dict[str, Any] = dict( - env=env, - stdin=subprocess.DEVNULL, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - if os.name == "posix": - popen_kwargs["start_new_session"] = True - else: - detached = getattr(subprocess, "DETACHED_PROCESS", 0) - new_group = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0) - popen_kwargs["creationflags"] = detached | new_group - proc = subprocess.Popen(cmd, **popen_kwargs) + cmd += ["--properties-file", conf_file] try: - # Wait for the server to write the discovery file (which records the actual bound port) - # and start accepting connections. + result = subprocess.run( + cmd, + env=env, + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + timeout=120, + ) + if result.returncode != 0: + stale_pid = _read_pid(pid_file) + if stale_pid is not None and _pid_alive(stale_pid): + # spark-daemon.sh refuses to start while its pid file points at a live + # process -- here a server this client just rejected as not reusable + # (e.g. after a Spark upgrade). + reason = ( + "a local Connect server that is not reusable by this client is already " + "running (pid {}); stop it with " + "`python -m pyspark.sql.connect.local_server --stop`".format(stale_pid) + ) + else: + output = (result.stderr or "") + (result.stdout or "") + last_line = output.strip().splitlines()[-1] if output.strip() else "" + reason = "start-connect-server.sh exited with code {}: {}".format( + result.returncode, last_line + ) + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={"reason": reason}, + ) + + # The script has daemonized the server; wait for it to accept connections. deadline = time.time() + 120 while time.time() < deadline: - exit_code = proc.poll() - if exit_code is not None: - # Our daemon died. Another process may have published a usable server in the - # meantime (e.g. a port race), so prefer reconnecting to it over failing. - endpoint = _reuse_from_discovery() - if endpoint is not None: - return endpoint - raise PySparkRuntimeError( - errorClass="LOCAL_CONNECT_SERVER_START_FAILED", - messageParameters={ - "reason": "the server process exited with code {}".format(exit_code) - }, - ) - disc = _read_discovery() - if disc is not None and disc["pid"] == proc.pid and disc["token"] == token: + pid = _read_pid(pid_file) + if pid is not None: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.settimeout(0.5) - if sock.connect_ex((disc["host"], disc["port"])) == 0: + if sock.connect_ex(("localhost", port)) == 0: + _write_discovery(discovery_path, "localhost", port, token, pid, __version__) os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = token - return "sc://{}:{}".format(disc["host"], disc["port"]) + return "sc://localhost:{}".format(port) + if not _pid_alive(pid): + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={ + "reason": "the server exited during start-up; see logs under " + "{}".format(log_dir) + }, + ) time.sleep(0.25) - # Timed out. The daemon may still be inside getOrCreate() -- i.e. before it has wired up - # its own SIGTERM handler -- and it has already spawned a child JVM. Signal the whole - # process group (the daemon is a session leader via start_new_session) and escalate to - # SIGKILL, so the JVM is reaped rather than orphaned. - _terminate_server(proc) + # Timed out: best-effort stop of the server we just started, then fail. + pid = _read_pid(pid_file) + if pid is not None: + _signal_server(pid, signal.SIGTERM) raise PySparkRuntimeError( errorClass="LOCAL_CONNECT_SERVER_START_FAILED", - messageParameters={"reason": "the server did not become ready within 120 seconds"}, + messageParameters={ + "reason": "the server did not become ready within 120 seconds; see logs " + "under {}".format(log_dir) + }, ) finally: if conf_file is not None: @@ -446,12 +420,21 @@ def reuse_or_start_local_connect_server(master: str, opts: Dict[str, Any]) -> st ``SparkSession._start_connect_server`` and is only reached for a ``local`` master when ``spark.local.connect.reuse`` / ``SPARK_LOCAL_CONNECT_REUSE`` is set. """ + if os.name != "posix": + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={ + "reason": "spark.local.connect.reuse relies on the POSIX scripts under sbin/; " + "on this platform start a server manually (sbin/start-connect-server.sh) and " + 'connect with .remote("sc://...")' + }, + ) # Fast path: reuse an already-running server without taking the cross-process lock. endpoint = _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 + # Without it, racing callers may each start a server; the winner writes the discovery file # and the others reconnect to it. with _start_lock(): endpoint = _reuse_from_discovery() @@ -472,158 +455,33 @@ def stop_local_connect_server() -> bool: stopped = False if disc is not None and disc["pid"] != os.getpid(): stopped = _signal_server(disc["pid"], signal.SIGTERM) - try: - os.remove(_discovery_path()) - except OSError: - pass + # Also drop the spark-daemon.sh pid file so a fresh start is never refused on its account. + for path in (_discovery_path(), _pid_file(_runtime_dir())): + try: + os.remove(path) + except OSError: + pass return stopped -# -- the daemon itself, run as `python -m pyspark.sql.connect.local_server` ----------------------- - - -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(description="Persistent local Spark Connect server.") - parser.add_argument( - "--stop", action="store_true", help="stop the recorded running server, if any, and exit" - ) - parser.add_argument("--master", default="local[*]") - parser.add_argument("--port", type=int, default=15002) - parser.add_argument( - "--discovery", - default=None, - help="path of the discovery file; defaults to $SPARK_LOCAL_CONNECT_DISCOVERY " - "or ~/.spark/connect-local.json", + parser = argparse.ArgumentParser( + description="Manage the persistent local Spark Connect server used by the opt-in " + "spark.local.connect.reuse path. The server itself is started on demand through " + "sbin/start-connect-server.sh." ) parser.add_argument( - "--conf-file", - default=None, - help="path to a JSON file of extra SparkConf entries to seed the server with", + "--stop", action="store_true", help="stop the recorded running server, if any" ) - parser.add_argument( - "--idle-timeout", - type=float, - default=1800.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() - if args.discovery: - # Everything below resolves the path via _discovery_path (as do the client-side helpers, - # e.g. under --stop), so map an explicit --discovery onto the override it honors. - os.environ["SPARK_LOCAL_CONNECT_DISCOVERY"] = args.discovery - discovery = _discovery_path() - - if args.stop: - if stop_local_connect_server(): - print("Stopped the persistent local Spark Connect server.") - else: - print("No running persistent local Spark Connect server found.") - return - - # 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) - - # The launcher passes the auth token through the environment (argv would leak it to `ps`); - # generate one for manual runs. The child JVM inherits the env var, and the Connect server - # falls back to it when the `spark.connect.authenticate.token` conf is not set. - token = os.environ.get("SPARK_CONNECT_AUTHENTICATE_TOKEN") - if not token: - token = str(uuid.uuid4()) - os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = 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") - ) - spark = builder.getOrCreate() - - bound_port = _bound_port(spark, args.port) - _write_discovery(discovery, "localhost", bound_port, token, __version__) - print( - "SPARK-CONNECT-LOCAL-SERVER READY port={} pid={}".format(bound_port, os.getpid()), - flush=True, - ) - - stop = False - - def _handle(_signum: int, _frame: Any) -> None: - nonlocal stop - stop = True - - signal.signal(signal.SIGTERM, _handle) - signal.signal(signal.SIGINT, _handle) - - idle_timeout = args.idle_timeout - poll_interval = max(1.0, args.poll_interval) - last_active = time.monotonic() - last_poll = 0.0 - try: - # Sleep in short slices so a SIGTERM is observed promptly regardless of the poll interval. - while not stop: - time.sleep(1.0) - if idle_timeout <= 0: - continue - now = time.monotonic() - if now - last_poll < poll_interval: - continue - last_poll = now - try: - active = _has_active_sessions(spark) - except Exception: - active = True # fail open: never reap a server we cannot inspect - if active: - last_active = now - elif now - last_active > idle_timeout: - break - finally: - _remove_discovery_if_ours(discovery) - spark.stop() + if not args.stop: + parser.print_help(sys.stderr) + sys.exit(2) + if stop_local_connect_server(): + print("Stopped the persistent local Spark Connect server.") + else: + print("No running persistent local Spark Connect server found.") if __name__ == "__main__": diff --git a/python/pyspark/sql/session.py b/python/pyspark/sql/session.py index 2466e75da693d..e9398bd512782 100644 --- a/python/pyspark/sql/session.py +++ b/python/pyspark/sql/session.py @@ -539,12 +539,9 @@ def getOrCreate(self) -> "SparkSession": # one on the first run) instead of booting a fresh in-process server # every process. See `pyspark.sql.connect.local_server`. url = reuse_or_start_local_connect_server(url, opts) - for k in ( - "spark.local.connect.reuse", - "spark.local.connect.server.port", - "spark.local.connect.server.idleTimeout", - ): - opts.pop(k, None) + for k in list(opts): + if k.startswith("spark.local.connect."): + opts.pop(k) elif url.startswith("local") or ( is_api_mode_connect and not url.startswith("sc://") ): diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server.py b/python/pyspark/sql/tests/connect/test_connect_local_server.py index 7793add5ad945..787c539f2e7e6 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server.py @@ -63,7 +63,7 @@ def tearDown(self) -> None: if disc is not None and disc["pid"] != os.getpid(): port = disc["port"] local_server.stop_local_connect_server() - # stop_local_connect_server only signals the daemon and returns; wait for the JVM + # stop_local_connect_server only signals the server and returns; wait for the JVM # to actually release the port so the next test starts from a clean slate. self._wait_port_closed(disc["host"], port) finally: @@ -154,7 +154,7 @@ def test_pid_probe_is_skipped_on_windows(self) -> None: def test_stop_when_no_server_is_safe(self) -> None: self.assertFalse(local_server.stop_local_connect_server()) - def test_stop_signals_recorded_server_group(self) -> None: + def test_stop_signals_recorded_server(self) -> None: from unittest import mock self._write_discovery(pid=12345) @@ -181,58 +181,57 @@ def test_stop_cli_reports_when_no_server(self) -> None: self.assertEqual(result.returncode, 0, result.stderr) self.assertIn("No running persistent local Spark Connect server", result.stdout) - @unittest.skipUnless(os.name == "posix", "process groups are POSIX-only") - def test_signal_group_kills_only_session_leaders(self) -> None: - """A pid that is not a session leader is signalled alone, never via its group. - - The daemon is always launched as a session leader (its group id equals its pid). A - recorded pid whose group id differs was recycled by an unrelated process -- for example - via a stale discovery file -- and group-killing it could take down the caller's own - process group. - """ + def test_reuse_or_start_requires_posix(self) -> None: + """The reuse path depends on the sbin shell scripts and must refuse to run elsewhere.""" from unittest import mock - sleeper = [sys.executable, "-c", "import time; time.sleep(60)"] - # Same process group as this test (not a leader): must fall back to a plain kill. - child = subprocess.Popen(sleeper) - try: - self.assertNotEqual(os.getpgid(child.pid), child.pid) - with mock.patch("os.killpg") as killpg, mock.patch("os.kill") as kill: - self.assertTrue(local_server._signal_server(child.pid, signal.SIGTERM)) - killpg.assert_not_called() - kill.assert_called_once_with(child.pid, signal.SIGTERM) - finally: - child.kill() - child.communicate() - # A session leader, like the real daemon: the whole group is signalled. - leader = subprocess.Popen(sleeper, start_new_session=True) - try: - self.assertEqual(os.getpgid(leader.pid), leader.pid) - with mock.patch("os.killpg") as killpg, mock.patch("os.kill") as kill: - self.assertTrue(local_server._signal_server(leader.pid, signal.SIGTERM)) - killpg.assert_called_once_with(leader.pid, signal.SIGTERM) - kill.assert_not_called() - finally: - leader.kill() - leader.communicate() + from pyspark.errors import PySparkRuntimeError + + with mock.patch.object(os, "name", "nt"): + with self.assertRaises(PySparkRuntimeError) as ctx: + local_server.reuse_or_start_local_connect_server("local[2]", {}) + self.assertIn("POSIX", str(ctx.exception)) def test_reuse_from_discovery_none_when_absent(self) -> None: self.assertIsNone(local_server._reuse_from_discovery()) - def test_local_port_available(self) -> None: + def test_pick_port_prefers_configured_and_falls_back(self) -> None: + """Outside SPARK_TESTING, the configured port is used unless it is already taken.""" + from unittest import mock + + env = {k: v for k, v in os.environ.items() if k != "SPARK_TESTING"} listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) try: listener.bind(("localhost", 0)) listener.listen(1) taken = listener.getsockname()[1] - self.assertFalse(local_server._port_available(taken)) + with mock.patch.dict(os.environ, env, clear=True): + # A taken configured port falls back to an OS-assigned one. + self.assertNotEqual( + local_server._pick_port({"spark.local.connect.server.port": taken}), taken + ) finally: listener.close() - # The port is free again once the listener is closed. - self.assertTrue(local_server._port_available(taken)) + # A free configured port is honored as-is (the listener above just released it). + with mock.patch.dict(os.environ, env, clear=True): + self.assertEqual( + local_server._pick_port({"spark.local.connect.server.port": taken}), taken + ) + + def test_write_seed_properties_perms_and_format(self) -> None: + """The seed --properties-file is 0600 and holds key=value lines.""" + path = local_server._write_seed_properties( + {"spark.sql.warehouse.dir": "/tmp/wh"}, self._tmpdir + ) + try: + self.assertEqual(os.stat(path).st_mode & 0o777, 0o600) + with open(path, "r") as f: + self.assertIn("spark.sql.warehouse.dir=/tmp/wh", f.read()) + finally: + os.remove(path) def test_server_conf_seeds_user_confs_and_drops_control_keys(self) -> None: - """_seed_conf keeps user startup confs but not keys the daemon controls.""" + """_seed_conf keeps user startup confs but not keys the launcher controls.""" opts = { "spark.sql.warehouse.dir": "/tmp/wh", "spark.jars.packages": "org.example:lib:1.0", @@ -242,7 +241,7 @@ def test_server_conf_seeds_user_confs_and_drops_control_keys(self) -> None: "spark.connect.grpc.binding.port": "15002", "spark.local.connect.reuse": "true", "spark.local.connect.server.port": "15002", - "spark.local.connect.server.idleTimeout": "60", + "spark.local.connect.future.knob": "x", # the whole prefix is reserved and dropped } conf = local_server._seed_conf(opts) self.assertEqual(conf.get("spark.sql.warehouse.dir"), "/tmp/wh") @@ -254,7 +253,7 @@ def test_server_conf_seeds_user_confs_and_drops_control_keys(self) -> None: "spark.connect.grpc.binding.port", "spark.local.connect.reuse", "spark.local.connect.server.port", - "spark.local.connect.server.idleTimeout", + "spark.local.connect.future.knob", ): self.assertNotIn(dropped, conf) @@ -268,21 +267,7 @@ def test_start_lock_roundtrip(self) -> None: return self.assertTrue(os.path.exists(self._discovery + ".lock")) - @unittest.skipUnless(os.name == "posix", "the removal guard needs fcntl") - def test_remove_discovery_skipped_while_start_lock_held(self) -> None: - """Daemon shutdown must not delete the discovery file while a launcher holds the lock. - - A held start-up lock means another process is publishing a new server; removing the file - in that window could delete the newcomer's entry (see _remove_discovery_if_ours). - """ - self._write_discovery(pid=os.getpid()) - with local_server._start_lock(): - local_server._remove_discovery_if_ours(self._discovery) - self.assertTrue(os.path.exists(self._discovery)) - local_server._remove_discovery_if_ours(self._discovery) - self.assertFalse(os.path.exists(self._discovery)) - - # -- end-to-end: start a real detached server, reconnect to it, verify isolation -------------- + # -- end-to-end: start a real server via sbin scripts, reconnect, verify isolation ------------ def _release(self, session) -> None: """Close one client session without stopping the shared server.""" @@ -306,6 +291,7 @@ def _wait_port_closed(self, host, port, timeout=30) -> bool: time.sleep(0.5) return False + @unittest.skipUnless(os.name == "posix", "the reuse path relies on the POSIX sbin scripts") def test_builder_remote_local_uses_reuse_flag(self) -> None: spark = None try: @@ -324,6 +310,7 @@ def test_builder_remote_local_uses_reuse_flag(self) -> None: if spark is not None: spark.stop() + @unittest.skipUnless(os.name == "posix", "the reuse path relies on the POSIX sbin scripts") def test_concurrent_startup_reuses_one_server(self) -> None: script = textwrap.dedent(""" import json @@ -376,8 +363,10 @@ def test_concurrent_startup_reuses_one_server(self) -> None: self.assertEqual(len({o["pid"] for o in outputs}), 1) self.assertEqual(len({o["port"] for o in outputs}), 1) + @unittest.skipUnless(os.name == "posix", "the reuse path relies on the POSIX sbin scripts") def test_start_reuse_and_session_isolation(self) -> None: - # First call starts a detached persistent server and records it in the discovery file. + # First call starts a persistent server via sbin scripts and records it in the discovery + # file. endpoint = local_server.reuse_or_start_local_connect_server("local[2]", {}) self.assertTrue(endpoint.startswith("sc://localhost:")) @@ -413,17 +402,17 @@ def test_start_reuse_and_session_isolation(self) -> None: # Stopping signals the server and removes the discovery file. self.assertTrue(local_server.stop_local_connect_server()) self.assertIsNone(local_server._read_discovery()) - # The server should stop accepting connections shortly afterwards. (We check the port rather - # than the pid: the detached server is a child of this long-lived test process, so once it - # exits it lingers as an unreaped zombie for which os.kill(pid, 0) still succeeds.) + # The server should stop accepting connections shortly afterwards. (We check the port + # rather than the pid, which can linger briefly while the JVM shuts down.) _, _, hostport = endpoint.partition("sc://") host, _, port = hostport.partition(":") self.assertTrue( self._wait_port_closed(host, port), "server port {} still open after stop".format(port) ) + @unittest.skipUnless(os.name == "posix", "the reuse path relies on the POSIX sbin scripts") def test_start_seeds_static_conf_on_the_server(self) -> None: - """A start-up conf passed by the first caller reaches the daemon's SparkConf.""" + """A start-up conf passed by the first caller reaches the server's SparkConf.""" warehouse = os.path.join(self._tmpdir, "seeded-wh") opts = {"spark.sql.warehouse.dir": warehouse} endpoint = local_server.reuse_or_start_local_connect_server("local[2]", opts) @@ -437,168 +426,6 @@ def test_start_seeds_static_conf_on_the_server(self) -> None: self._release(spark) # tearDown stops the server and waits for the port to close. - @staticmethod - def _running_group_members(pgid: int) -> list: - """Pids in process group ``pgid`` that are still running, excluding zombies. - - Zombies count as gone: they are already dead, so nothing can leak. Signal-based checks - like ``killpg(pgid, 0)`` cannot be used instead, because a zombie keeps its group - signalable, and orphans reparented to a pid 1 that never reaps (e.g. minimal CI - containers) stay zombies indefinitely. - """ - out = subprocess.run( - ["ps", "-A", "-o", "pid=,pgid=,stat="], capture_output=True, text=True - ).stdout - members = [] - for line in out.splitlines(): - parts = line.split() - if len(parts) >= 3 and parts[1] == str(pgid) and not parts[2].startswith("Z"): - members.append(int(parts[0])) - return members - - def _wait_process_group_gone(self, pgid: int, timeout: float) -> bool: - """Wait until nothing in ``pgid`` is still running; True if that happened in time.""" - deadline = time.time() + timeout - while True: - if not self._running_group_members(pgid): - return True - if time.time() >= deadline: - return False - time.sleep(0.2) - - @unittest.skipUnless(os.name == "posix", "process-group reaping is POSIX-only") - def test_terminate_escalates_when_group_outlives_daemon(self) -> None: - """SIGKILL escalation covers group members that survive the daemon. - - Reproduces the CI failure mode of test_terminate_reaps_daemon_and_jvm: the daemon exits - on SIGTERM promptly while its JVM is still shutting down (or wedged) past the grace - period. _terminate_server must not return while the group has live members. - """ - if shutil.which("ps") is None: - self.skipTest("ps is needed to inspect process group members") - - # A session leader that exits on SIGTERM after spawning a SIGTERM-immune child into its - # process group -- the child stands in for a JVM that outlives the daemon. - leader_prog = textwrap.dedent(""" - import signal - import subprocess - import sys - import time - - child = subprocess.Popen( - [ - sys.executable, - "-c", - "import signal, time; " - "signal.signal(signal.SIGTERM, signal.SIG_IGN); " - "print('ready', flush=True); " - "time.sleep(600)", - ] - ) - print("pid:%d" % child.pid, flush=True) - signal.signal(signal.SIGTERM, lambda *_: sys.exit(0)) - time.sleep(600) - """) - proc = subprocess.Popen( - [sys.executable, "-c", leader_prog], - stdout=subprocess.PIPE, - start_new_session=True, - ) - pgid = os.getpgid(proc.pid) - try: - # One line comes from the leader (the child's pid) and one from the child itself, - # printed only after it has ignored SIGTERM -- reading both closes the start-up race. - lines = {proc.stdout.readline().strip() for _ in range(2)} - self.assertIn(b"ready", lines) - child_pid = int(next(ln for ln in lines if ln.startswith(b"pid:"))[4:]) - - local_server._terminate_server(proc) - - self.assertTrue( - self._wait_process_group_gone(pgid, timeout=10), - "process group survived _terminate_server", - ) - # The SIGTERM-immune child no longer runs, proving the SIGKILL escalation fired. - self.assertNotIn( - child_pid, - self._running_group_members(pgid), - "SIGTERM-immune child survived the SIGKILL escalation", - ) - finally: - proc.stdout.close() - try: - os.killpg(pgid, signal.SIGKILL) - except OSError: - pass - - @unittest.skipUnless(os.name == "posix", "process-group reaping is POSIX-only") - def test_terminate_reaps_daemon_and_jvm(self) -> None: - """_terminate_server kills the daemon *and* its child JVM. - - The test waits until the daemon has launched the JVM, then checks that the whole process - group exits. - """ - if shutil.which("pgrep") is None: - self.skipTest("pgrep is needed to detect the child JVM") - - cmd = [ - sys.executable, - "-m", - "pyspark.sql.connect.local_server", - "--master", - "local[2]", - "--port", - "0", # ephemeral, so a stray server elsewhere cannot interfere - "--discovery", - self._discovery, - "--idle-timeout", - "3600", - ] - env = dict(os.environ) - for var in ("SPARK_REMOTE", "SPARK_LOCAL_REMOTE", "SPARK_CONNECT_MODE_ENABLED"): - env.pop(var, None) - proc = subprocess.Popen( - cmd, - env=env, - stdin=subprocess.DEVNULL, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - start_new_session=True, # session/group leader, matching the production launch - ) - - def child_pids(): - try: - out = subprocess.check_output(["pgrep", "-P", str(proc.pid)]) - except subprocess.CalledProcessError: - return [] - return [int(p) for p in out.split()] - - pgid = os.getpgid(proc.pid) - try: - # Wait until the daemon has spawned its child JVM before terminating the group. - deadline = time.time() + 90 - kids = [] - while time.time() < deadline: - kids = child_pids() - if kids: - break - if proc.poll() is not None: - self.fail("daemon exited (code {}) before launching a JVM".format(proc.poll())) - time.sleep(0.5) - self.assertTrue(kids, "daemon never launched a JVM; cannot exercise the orphan case") - - local_server._terminate_server(proc) - - self.assertTrue( - self._wait_process_group_gone(pgid, timeout=30), - "process group survived _terminate_server", - ) - finally: - try: - os.killpg(pgid, signal.SIGKILL) - except OSError: - pass - if __name__ == "__main__": from pyspark.testing import main From f72f07f1f84efa33147db353d79ee1890e568896 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Wed, 8 Jul 2026 17:25:20 -0700 Subject: [PATCH 13/14] [SPARK-57787][CONNECT] Keep the reuse conf internal; document only the manual server workflow Per review discussion, spark.local.connect.reuse / SPARK_LOCAL_CONNECT_REUSE stays an internal, undocumented conf. The docs now cover only the explicit workflow (sbin/start-connect-server.sh + .remote("sc://...")) plus a note on what is and is not isolated between runs against a shared local server. Co-authored-by: Isaac --- docs/spark-connect-overview.md | 47 +++++----------------------------- 1 file changed, 7 insertions(+), 40 deletions(-) diff --git a/docs/spark-connect-overview.md b/docs/spark-connect-overview.md index 72896b0bfa83f..a517c40193dca 100644 --- a/docs/spark-connect-overview.md +++ b/docs/spark-connect-overview.md @@ -291,11 +291,8 @@ PySpark boots a fresh in-process Spark Connect server in **every** process. Each JVM warmup, `SparkContext` construction, and Connect server boot -- which can take a few seconds and makes a quick edit/run loop feel slow. -There are two ways to amortize that cost across runs by reconnecting to a long-lived local server. - -### Start a server yourself and connect to it - -Start one persistent local Spark Connect server and point every run at it: +To amortize that cost across runs, start one persistent local Spark Connect server and point +every run at it: ```bash # Start once; it stays up across runs. @@ -308,41 +305,11 @@ python -c 'from pyspark.sql import SparkSession; SparkSession.builder.remote("sc $SPARK_HOME/sbin/stop-connect-server.sh ``` -### Let PySpark manage the server (opt-in) - -If you would rather keep your code as `SparkSession.builder.remote("local[*]").getOrCreate()` and not -manage a server by hand, enable the opt-in reuse path. The first run starts one persistent local -Connect server through the standard `sbin/start-connect-server.sh` script -- the same daemon you -would start by hand above -- and records it in a discovery file; later runs reconnect to it in a -fraction of a second: - -```bash -export SPARK_LOCAL_CONNECT_REUSE=1 # or .config("spark.local.connect.reuse", "true") -python script.py # 1st run: starts a persistent server (cold start, once) -python script.py # 2nd+ run: reconnects to it (sub-second) -``` - -This is **off by default**; nothing changes unless you opt in. A few details: - -- Each run connects as its own Connect session, so session-local state -- temp views, runtime SQL - configurations, and (with artifact isolation, which stays on) session artifacts -- is fresh on - every run and never leaks between runs. State backed by the shared `SparkContext` (the persistent - catalog/warehouse, global temp views, and cached datasets) *is* shared across runs, so namespace - per-run databases or clear that state yourself if your runs must be fully isolated. -- The server listens on port `15002` by default and authenticates with a token written, together - with the host, port, pid and Spark version, to `~/.spark/connect-local.json` (mode `0600`). - Set `SPARK_LOCAL_CONNECT_DISCOVERY` to relocate that file. Reuse is refused (and a fresh server - started) if the recorded process is gone, the port is closed, or the Spark version differs. On - Unix-like systems, PySpark uses a file lock around first startup. On platforms without `fcntl`, - concurrent startups can race; callers that lose the race reconnect to the server recorded in the - discovery file. -- The server is a regular `spark-daemon.sh`-managed process and runs until you stop it with - `python -m pyspark.sql.connect.local_server --stop` (or `sbin/stop-connect-server.sh`, or by - terminating the `pid` recorded in the discovery file). Its pid file and logs live next to the - discovery file (`~/.spark` and `~/.spark/logs` by default), which is where to look if a start-up - fails. After a Spark upgrade the old server is rejected on its recorded version; stop it the same - way. Because this path drives the `sbin/` shell scripts, it is unavailable on Windows -- there, - start a server manually and connect to it with `.remote("sc://...")`. +Each run connects as its own Connect session, so session-local state -- temp views, runtime SQL +configurations, and session artifacts -- is fresh on every run and never leaks between runs. State +backed by the shared `SparkContext` (the persistent catalog/warehouse, global temp views, and +cached datasets) *is* shared across runs, so namespace per-run databases or clear that state +yourself if your runs must be fully isolated. ## Use Spark Connect in standalone applications From 61d5c950890a536a859e033d36d64f50efad6c1a Mon Sep 17 00:00:00 2001 From: ericm-db Date: Thu, 9 Jul 2026 10:11:05 -0700 Subject: [PATCH 14/14] [SPARK-57787][CONNECT] Fix ruff format in local_server.py Co-authored-by: Isaac --- python/pyspark/sql/connect/local_server.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/pyspark/sql/connect/local_server.py b/python/pyspark/sql/connect/local_server.py index 9019750a06420..529be5cf77f59 100644 --- a/python/pyspark/sql/connect/local_server.py +++ b/python/pyspark/sql/connect/local_server.py @@ -388,8 +388,9 @@ def _launch_server(master: str, opts: Dict[str, Any]) -> str: raise PySparkRuntimeError( errorClass="LOCAL_CONNECT_SERVER_START_FAILED", messageParameters={ - "reason": "the server exited during start-up; see logs under " - "{}".format(log_dir) + "reason": "the server exited during start-up; see logs under {}".format( + log_dir + ) }, ) time.sleep(0.25)