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..a517c40193dca 100644
--- a/docs/spark-connect-overview.md
+++ b/docs/spark-connect-overview.md
@@ -277,6 +277,40 @@ 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
+
+```python
+from pyspark.sql import SparkSession
+spark = SparkSession.builder.remote("local[*]").getOrCreate()
+```
+
+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.
+
+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.
+$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
+```
+
+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
@@ -371,7 +405,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/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/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..529be5cf77f59
--- /dev/null
+++ b/python/pyspark/sql/connect/local_server.py
@@ -0,0 +1,489 @@
+#
+# 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.
+#
+
+"""
+Opt-in reuse of a persistent local Spark Connect server (``spark.local.connect.reuse`` /
+``SPARK_LOCAL_CONNECT_REUSE``).
+
+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, 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
+import contextlib
+import json
+import os
+import signal
+import socket
+import subprocess
+import sys
+import tempfile
+import time
+import uuid
+from typing import Any, Dict, Iterator, Optional
+
+from pyspark.errors import PySparkRuntimeError
+
+_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:
+ """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 _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.
+
+ 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, 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:
+ os.makedirs(parent, exist_ok=True)
+ payload = {
+ "host": host,
+ "port": port,
+ "token": token,
+ "pid": pid,
+ "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)
+
+
+# -- deciding between reusing the recorded server and starting a fresh one ------------------------
+
+
+def _pid_alive(pid: int) -> bool:
+ """Whether ``pid`` exists (POSIX only). A process we cannot signal counts as alive."""
+ try:
+ os.kill(pid, 0)
+ except ProcessLookupError:
+ return False
+ except OSError:
+ pass
+ return True
+
+
+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" 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:
+ 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 server, and
+ the losers reconnect to the one 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
+
+
+# -- 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 port
+ except OSError:
+ return free_port()
+
+
+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 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 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 _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``.
+
+ 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.
+ """
+ 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 _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))
+
+
+def _read_pid(path: str) -> Optional[int]:
+ try:
+ 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.kill(pid, sig)
+ return True
+ except OSError:
+ return False
+
+
+def _launch_server(master: str, opts: Dict[str, Any]) -> str:
+ """Start a persistent local Connect server and wait until it is reachable.
+
+ 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.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. 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)
+
+ 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 = [
+ script,
+ "--master",
+ master,
+ "--conf",
+ "spark.connect.grpc.binding.port={}".format(port),
+ ]
+ if conf_file is not None:
+ cmd += ["--properties-file", conf_file]
+
+ try:
+ 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:
+ 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(("localhost", port)) == 0:
+ _write_discovery(discovery_path, "localhost", port, token, pid, __version__)
+ os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = token
+ 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: 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; see logs "
+ "under {}".format(log_dir)
+ },
+ )
+ 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.
+ """
+ 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 server; 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)
+ # 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
+
+
+def main() -> None:
+ 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(
+ "--stop", action="store_true", help="stop the recorded running server, if any"
+ )
+ args = parser.parse_args()
+
+ 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__":
+ main()
diff --git a/python/pyspark/sql/session.py b/python/pyspark/sql/session.py
index c36260b9d13ea..e9398bd512782 100644
--- a/python/pyspark/sql/session.py
+++ b/python/pyspark/sql/session.py
@@ -523,7 +523,26 @@ 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:
+ 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 `pyspark.sql.connect.local_server`.
+ url = reuse_or_start_local_connect_server(url, opts)
+ 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://")
):
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..787c539f2e7e6
--- /dev/null
+++ b/python/pyspark/sql/tests/connect/test_connect_local_server.py
@@ -0,0 +1,433 @@
+#
+# 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 shutil
+import signal
+import socket
+import subprocess
+import sys
+import tempfile
+import textwrap
+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 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__
+
+
+@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 = 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 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:
+ for k, v in self._saved_env.items():
+ if v is None:
+ os.environ.pop(k, None)
+ else:
+ os.environ[k] = v
+ # 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) ----------------------------------------
+
+ def test_discovery_path_honors_override(self) -> None:
+ self.assertEqual(local_server._discovery_path(), self._discovery)
+ os.environ.pop("SPARK_LOCAL_CONNECT_DISCOVERY")
+ self.assertTrue(
+ local_server._discovery_path().endswith(os.path.join(".spark", "connect-local.json"))
+ )
+
+ def test_read_discovery_missing_or_malformed(self) -> None:
+ self.assertIsNone(local_server._read_discovery())
+ with open(self._discovery, "w") as f:
+ f.write("not json")
+ self.assertIsNone(local_server._read_discovery())
+ with open(self._discovery, "w") as f:
+ json.dump({"host": "localhost"}, f) # missing required keys
+ 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 = {
+ "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(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(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.
+ 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(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(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(local_server.stop_local_connect_server())
+
+ def test_stop_signals_recorded_server(self) -> None:
+ from unittest import mock
+
+ self._write_discovery(pid=12345)
+ calls = []
+
+ def fake_signal(pid, sig):
+ calls.append((pid, sig))
+ return True
+
+ 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)
+
+ 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
+
+ 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_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]
+ 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()
+ # 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 launcher 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.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")
+ 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.future.knob",
+ ):
+ self.assertNotIn(dropped, conf)
+
+ def test_start_lock_roundtrip(self) -> None:
+ """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"))
+
+ # -- 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."""
+ try:
+ session.client.release_session()
+ except Exception:
+ pass
+ try:
+ session.client.close()
+ 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
+
+ @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:
+ spark = (
+ PySparkSession.builder.remote("local[2]")
+ .config("spark.local.connect.reuse", "true")
+ .getOrCreate()
+ )
+ self.assertEqual(spark.range(2).count(), 2)
+
+ disc = local_server._read_discovery()
+ self.assertIsNotNone(disc)
+ self.assertEqual(disc["spark_version"], __version__)
+ self.assertNotEqual(disc["pid"], os.getpid())
+ finally:
+ 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
+ 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)
+
+ @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 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:"))
+
+ 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"])
+ first_pid = disc["pid"]
+
+ s1 = s2 = None
+ try:
+ # A second call reuses the running server: same endpoint, no new process spawned.
+ endpoint2 = local_server.reuse_or_start_local_connect_server("local[2]", {})
+ self.assertEqual(endpoint2, endpoint)
+ 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()
+ 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(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, 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 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)
+ spark = None
+ try:
+ spark = RemoteSparkSession.builder.remote(endpoint).create()
+ # 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:
+ self._release(spark)
+ # tearDown stops the server and waits for the port to close.
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()