dtenedor commented on code in PR #56907: URL: https://github.com/apache/spark/pull/56907#discussion_r3539980768
########## python/pyspark/sql/tests/connect/test_connect_local_server.py: ########## @@ -0,0 +1,353 @@ +# +# 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 sys +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__ + + [email protected]( + 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(): + 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: + 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(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()) + + 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() Review Comment: Windows has no file lock: `_acquire_local_connect_start_lock` returns `None` without `fcntl`; relies on reconnect-the-winner fallback. Acceptable if documented; races are possible on Windows. ########## python/pyspark/sql/connect/session.py: ########## @@ -1247,6 +1249,347 @@ 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``, 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 _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. + + 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: Review Comment: This signals only the daemon PID, not the process group. Unlike the timeout path, stop uses `os.kill(pid, SIGTERM)` only. If called while the daemon is still inside `getOrCreate()` (before its signal handler is installed), the JVM could be orphaned — same class of bug that was fixed for timeout. Unlikely in normal use (stop is a dev helper after the server is up), but inconsistent with `_terminate_local_connect_server`. Consider reusing the group-kill helper. ########## python/pyspark/sql/tests/connect/test_connect_local_server.py: ########## @@ -0,0 +1,353 @@ +# +# 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 sys +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__ + + [email protected]( + 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): Review Comment: Could we also add an automated concurrent-startup test? And one integration test through `SparkSession.builder.remote("local[*]").getOrCreate()` with the env var set? ########## docs/spark-connect-overview.md: ########## @@ -277,6 +277,65 @@ The connection may also be programmatically created using _SparkSession#builder_ </div> </div> +## 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 Review Comment: On version mismatch, a new server starts on an ephemeral port and writes discovery. The old server on 15002 keeps running until idle timeout. Over repeated upgrades, multiple idle JVMs could accumulate. Docs could mention killing the old pid from a stale discovery file. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
