diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index 898a3831f484b..b90e28a20d005 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -1179,6 +1179,7 @@ def __hash__(self): "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_local_server_pool", "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 2f1c0836c3837..64e184cf723bb 100644 --- a/docs/spark-connect-overview.md +++ b/docs/spark-connect-overview.md @@ -343,6 +343,37 @@ backed by the shared `SparkContext` (the persistent catalog/warehouse, global te cached datasets) *is* shared across runs, so namespace per-run databases or clear that state yourself if your runs must be fully isolated. +### Fully isolated runs with a pool of single-use servers + +If your runs must be fully isolated from each other but you still want to skip the per-run +startup cost, a second experimental opt-in keeps a small pool of booted servers that have +never been assigned to an application run instead of one shared one. With +`SPARK_LOCAL_CONNECT_POOL=1` set (or +`spark.local.connect.pool=true` on the builder), each run *claims* a fresh server from the pool, +a replacement is booted in the background, and the claimed server is torn down when the run's +session stops -- no server ever serves two runs, so runs are as isolated as with the default +in-process server: + +```bash +export SPARK_LOCAL_CONNECT_POOL=1 + +# Each run claims a pre-booted server not previously assigned to another run. +python -c 'from pyspark.sql import SparkSession; SparkSession.builder.remote("local[*]").getOrCreate()' + +# Force-stop all pool servers and start over from a clean slate. +python -m pyspark.sql.connect.local_server_pool --purge +``` + +The trade-off relative to the reuse mode above is memory: `spark.local.connect.pool.size` +(default 2) idle servers stay resident while you iterate, several hundred MB each. Idle pool +servers warm themselves up with a fixed set of synthetic queries (disable with +`SPARK_LOCAL_CONNECT_POOL_WARMUP=0`), so the first real query of a run is fast too, and they +shut down on their own after sitting unclaimed for `SPARK_LOCAL_CONNECT_POOL_IDLE_TIMEOUT` +seconds (default 1800), so an idle machine drains back to zero servers. Pool servers are only +handed to runs whose master, startup configurations, working directory, and Python environment +match the ones they were booted with; runs that differ in any of these boot their own pool +members. If both this and `spark.local.connect.reuse` are set, the pool takes precedence. + ## 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 4f06d884aff49..30f4b1939fe3e 100644 --- a/python/pyspark/sql/connect/local_server.py +++ b/python/pyspark/sql/connect/local_server.py @@ -75,6 +75,22 @@ def _pid_alive(pid: int) -> bool: return True +def runtime_dir() -> str: + """Return the private per-user directory holding local-server state.""" + path = os.path.join(tempfile.gettempdir(), "spark-connect-{}".format(getpass.getuser())) + try: + # exist_ok also covers two first runs racing to create the directory; chmod + # re-asserts 0700 and fails if another user owns the path. + os.makedirs(path, mode=0o700, exist_ok=True) + os.chmod(path, 0o700) + except OSError as e: + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_RUNTIME_DIR_UNAVAILABLE", + messageParameters={"path": path}, + ) from e + return path + + class Discovery: """Reads and writes the discovery file recording the persistent local server. @@ -82,26 +98,11 @@ class Discovery: ``SPARK_LOCAL_CONNECT_DISCOVERY`` points; the daemon's pid file and logs sit next to it. """ - @staticmethod - def _runtime_dir() -> str: - path = os.path.join(tempfile.gettempdir(), "spark-connect-{}".format(getpass.getuser())) - try: - # exist_ok also covers two first runs racing to create the directory; chmod - # re-asserts 0700 and fails if another user owns the path. - os.makedirs(path, mode=0o700, exist_ok=True) - os.chmod(path, 0o700) - except OSError as e: - raise PySparkRuntimeError( - errorClass="LOCAL_CONNECT_RUNTIME_DIR_UNAVAILABLE", - messageParameters={"path": path}, - ) from e - return path - def __init__(self, path: Optional[str] = None): self.path = os.path.abspath( path or os.environ.get("SPARK_LOCAL_CONNECT_DISCOVERY") - or os.path.join(self._runtime_dir(), "connect-local.json") + or os.path.join(runtime_dir(), "connect-local.json") ) self._lock_file: Optional[TextIO] = None @@ -219,12 +220,33 @@ def is_reusable(self) -> bool: def reuse_or_start(self, master: str, opts: Dict[str, Any]) -> str: if not self.is_reusable(): - ServerLauncher(master, opts, self._discovery).launch() - self._reload() + self.start(master, opts) assert self.token is not None os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = self.token return self.url + def start( + self, + master: str, + opts: Dict[str, Any], + *, + use_ephemeral_port: bool = False, + seed_conf: Optional[Dict[str, Any]] = None, + ) -> None: + """Start this server and reload its discovery record. + + Callers starting isolated daemons can request an ephemeral port and provide a + precomputed startup configuration while sharing the standard launch path. + """ + ServerLauncher( + master, + opts, + self._discovery, + use_ephemeral_port=use_ephemeral_port, + seed_conf=seed_conf, + ).launch() + self._reload() + def stop(self) -> bool: stopped = False if self.pid is not None: @@ -237,17 +259,47 @@ def stop(self) -> bool: return stopped +def startup_seed_conf(opts: Dict[str, Any]) -> Dict[str, Any]: + """Compute startup confs using the same merge as the in-process server path.""" + 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 + + class ServerLauncher: """Starts a persistent local server via ``sbin/start-connect-server.sh`` and waits until it accepts connections. Callers must hold a ``Discovery`` context. + + ``use_ephemeral_port`` and ``seed_conf`` allow other managed local servers to share this + launch path without duplicating its process and readiness handling. """ _READY_TIMEOUT = 120 - def __init__(self, master: str, opts: Dict[str, Any], discovery: Discovery): + def __init__( + self, + master: str, + opts: Dict[str, Any], + discovery: Discovery, + use_ephemeral_port: bool = False, + seed_conf: Optional[Dict[str, Any]] = None, + ): self._master = master self._opts = opts self._discovery = discovery + self._use_ephemeral_port = use_ephemeral_port + self._seed_override = seed_conf self._log_dir = os.path.join(discovery.directory, "logs") def launch(self) -> None: @@ -269,12 +321,12 @@ def _token(self) -> str: ) def _pick_port(self) -> int: - """Under SPARK_TESTING always use an OS-assigned free port so suites can run in - parallel; otherwise honor the configured/default port, falling back to a free one if - another process holds it. (A live stale server of ours also holds the port, but that - start fails later at spark-daemon.sh's pid-file check regardless of port.) The sbin - script cannot report an ephemeral port back, so the free port is picked and released - here, with a small race until the server binds it. + """Use an OS-assigned free port when requested or under SPARK_TESTING so suites can + run in parallel. Otherwise honor the configured/default port, falling back to a free + one if another process holds it. (A live stale server of ours also holds the port, but + that start fails later at spark-daemon.sh's pid-file check regardless of port.) The + sbin script cannot report an ephemeral port back, so the free port is picked and + released here, with a small race until the server binds it. """ def free_port() -> int: @@ -282,7 +334,7 @@ def free_port() -> int: sock.bind(("localhost", 0)) return sock.getsockname()[1] - if "SPARK_TESTING" in os.environ: + if self._use_ephemeral_port or "SPARK_TESTING" in os.environ: return free_port() from pyspark.sql.connect.client import DefaultChannelBuilder @@ -303,20 +355,9 @@ def _seed_conf(self) -> Dict[str, Any]: opt-in keys. Only the run that starts the server can seed static confs; later runs find the JVM already warm. """ - 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(self._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 + if self._seed_override is not None: + return dict(self._seed_override) + return startup_seed_conf(self._opts) @contextlib.contextmanager def _seed_properties_file(self) -> Iterator[Optional[str]]: diff --git a/python/pyspark/sql/connect/local_server_pool.py b/python/pyspark/sql/connect/local_server_pool.py new file mode 100644 index 0000000000000..49235ca155f1e --- /dev/null +++ b/python/pyspark/sql/connect/local_server_pool.py @@ -0,0 +1,892 @@ +# +# 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 pool of pre-warmed, single-use local Spark Connect servers +(``spark.local.connect.pool`` / ``SPARK_LOCAL_CONNECT_POOL``). + +The reuse mode (``spark.local.connect.reuse``, see ``pyspark.sql.connect.local_server``) makes +local runs fast by sharing one long-lived server, at the price of state backed by the shared +``SparkContext`` (persistent catalog, global temp views, cached data) carrying across runs. +The pool keeps the speed without the sharing: it maintains a small set of booted servers that +have never been assigned to an application run, and +``SparkSession.builder.remote("local[*]").getOrCreate()`` *claims* one exclusively, spawns a +replacement in the background, and tears the claimed server down when the session stops or the +client exits. No server ever serves two application runs, so runs are as isolated from each +other as with the default in-process server -- at the cost of the idle servers' memory while +you iterate. If both opt-ins are set, the pool takes precedence. + +The pool lives in a ``pool`` subdirectory of the per-user runtime directory (override with +``SPARK_LOCAL_CONNECT_POOL_DIR``). A member is a set of files named by a random ````; +every access happens under the directory's ``.lock`` file lock, so readers always observe +complete states: + + pending-.json an in-flight launch (the attendant process booting the server) + conf-.json startup confs for that launch, read once by its attendant + server-.json a ready, unclaimed server (host/port/token/pid/version) + claimed--.json a server owned by the live client process + retired-.json a server being torn down; hard-killed if it hangs + member-/ the server's pid file and logs (spark-daemon.sh directories) + +Each launch runs an *attendant* (``python -m pyspark.sql.connect.local_server_pool --attend``), +a small detached process that boots the server through ``sbin/start-connect-server.sh``, +publishes its ``server-.json``, JIT-warms it with fixed synthetic queries in a bounded +child process (so the one real run the member serves starts with hot JIT and codegen caches; +disable with ``SPARK_LOCAL_CONNECT_POOL_WARMUP=0``), and then stays on as the member's +supervisor: it retires the server once it has sat unclaimed past the idle timeout +(``SPARK_LOCAL_CONNECT_POOL_IDLE_TIMEOUT`` seconds, default 1800), or once the client that +claimed it has died without releasing it. An idle machine therefore drains back to zero +servers without any further Spark run. A janitor pass on every acquire is the backstop for +members whose attendant itself died. + +Servers are only handed to runs they were built for: each member carries a fingerprint of its +master, seeded confs, working directory, and Python executable, and a run only claims members +whose fingerprint matches its own. ``python -m pyspark.sql.connect.local_server_pool --purge`` +force-stops every member and empties the pool directory. + +This mode is experimental. Everything but the opt-in itself -- the pool directory layout, the +attendant, and the ``--purge`` entry point -- is an internal detail that may change or move, +for example into a unified ``spark connect`` CLI. POSIX only, like the reuse mode. +""" + +import argparse +import atexit +import contextlib +import hashlib +import json +import os +import shutil +import signal +import socket +import subprocess +import sys +import time +import uuid +from typing import Any, Dict, List, Optional, Tuple + +from pyspark.errors import PySparkRuntimeError + +# How long one acquire may wait for a member to become ready. A cold launch takes at most +# 120s; this leaves room for one relaunch of a failed one. +_ACQUIRE_TIMEOUT = 180 +# A pending marker older than this belongs to a launch that hung; the janitor kills it. +# Deliberately above the local-server startup timeout so a slow-but-healthy launch is never shot. +_LAUNCH_TIMEOUT = 180 +# A retired server still alive this long after retirement is hard-killed; one that survives +# even that (e.g. not ours to signal) is dropped from tracking after the give-up age. +_RETIRE_KILL_AFTER = 30 +_RETIRE_GIVE_UP = 600 +# Unreferenced member directories older than this are removed. The age gate keeps the logs of +# a just-failed launch around long enough to be looked at. +_DIR_GC_AGE = 24 * 3600 + +_DEFAULT_POOL_SIZE = 2 +_DEFAULT_IDLE_TIMEOUT = 1800 + + +def _pool_size(opts: Dict[str, Any]) -> int: + """The number of warm or in-flight members to keep per fingerprint; at least one. + Malformed values fall back to the default rather than failing session creation over a + tuning knob. + """ + value = opts.get( + "spark.local.connect.pool.size", os.environ.get("SPARK_LOCAL_CONNECT_POOL_SIZE") + ) + try: + return max(1, int(value)) if value is not None else _DEFAULT_POOL_SIZE + except (TypeError, ValueError): + return _DEFAULT_POOL_SIZE + + +def _idle_timeout() -> int: + """Seconds an unclaimed member may sit before it is retired; 0 or negative disables idle + retirement. Read from the environment wherever reaping runs -- clients and attendants + alike -- so there is exactly one source of truth for it. + """ + try: + return int(os.environ["SPARK_LOCAL_CONNECT_POOL_IDLE_TIMEOUT"]) + except (KeyError, ValueError): + return _DEFAULT_IDLE_TIMEOUT + + +def _warmup_enabled() -> bool: + return os.environ.get("SPARK_LOCAL_CONNECT_POOL_WARMUP", "1").lower() not in ("0", "false") + + +def pool_fingerprint(master: str, seed_conf: Dict[str, Any]) -> str: + """The identity of a pool member: everything that shapes the server a run would have + booted for itself. A run only claims members whose fingerprint equals its own, so a + pre-booted JVM is never handed to a run it would not have produced. + + Besides the master and the seeded confs, this covers the working directory (unset + warehouse and Derby metastore locations resolve relative to it) and the Python executable + (the server runs Python UDFs with the interpreter environment it inherited from its + spawner). + """ + identity = [ + master, + sorted((str(k), str(v)) for k, v in seed_conf.items()), + os.getcwd(), + sys.executable, + os.environ.get("PYSPARK_PYTHON", ""), + ] + return hashlib.sha256(json.dumps(identity).encode("utf-8")).hexdigest()[:16] + + +def _pid_alive(pid: int) -> bool: + """Whether ``pid`` is running. A process we cannot signal counts as alive. Linux zombies + count as terminated: they remain signalable until their parent reaps them, but cannot own + or serve a pool member. POSIX only, like everything in this module. + """ + if pid <= 0: + return False + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except OSError: + pass + if sys.platform.startswith("linux"): + try: + with open(f"/proc/{pid}/stat", encoding="utf-8") as stat_file: + fields = stat_file.read().rpartition(")")[2].split() + except FileNotFoundError: + return False + except OSError: + pass + else: + if fields and fields[0] == "Z": + return False + return True + + +def _signal(pid: int, sig: int) -> bool: + """Best-effort signal; ``False`` when the process is already gone or not ours.""" + if pid <= 0: + return False + try: + os.kill(pid, sig) + return True + except OSError: + return False + + +class PoolMember: + """One published pool server, wrapping its ``server-.json`` record.""" + + def __init__(self, data: Dict[str, Any]): + self.data = data + # Set when this process claims the member; the path of its claimed--.json. + self.claim_path: Optional[str] = None + + @property + def pid(self) -> int: + return int(self.data["pid"]) + + @property + def token(self) -> str: + return self.data["token"] + + @property + def created(self) -> float: + return float(self.data.get("created", 0)) + + @property + def fingerprint(self) -> Optional[str]: + return self.data.get("fingerprint") + + @property + def url(self) -> str: + return f"sc://{self.data['host']}:{self.data['port']}" + + def is_usable(self) -> bool: + """Whether this member can serve a run: complete record, matching Spark version, live + process, and accepting connections.""" + from pyspark.version import __version__ + + if not all(k in self.data for k in ("host", "port", "token", "pid", "spark_version")): + return False + if self.data["spark_version"] != __version__ or not _pid_alive(self.pid): + return False + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(0.5) + return sock.connect_ex((self.data["host"], int(self.data["port"]))) == 0 + + +class PoolDirectory: + """Path layout, file access, and the cross-process lock of one pool directory. + + Used as a context manager that holds the directory's exclusive lock:: + + with PoolDirectory() as directory: + ...read and write state files... + + Acquire runs once per client process, so the simplicity of one exclusive lock for every + access beats any finer-grained scheme; the rename in ``ServerPool.claim`` is for tidiness, + not lock-free atomicity. The context is reentrant per instance in the sense that it can be + entered again after exiting, which the acquire loop does once per poll so that attendants + get their turn at publishing. + """ + + def __init__(self, path: Optional[str] = None): + if path is None: + path = os.environ.get("SPARK_LOCAL_CONNECT_POOL_DIR") + if path is None: + from pyspark.sql.connect.local_server import runtime_dir + + path = os.path.join(runtime_dir(), "pool") + self.path = os.path.abspath(path) + self._lock_fd: Optional[int] = None + + def __enter__(self) -> "PoolDirectory": + import fcntl + + os.makedirs(self.path, mode=0o700, exist_ok=True) + self._lock_fd = os.open(os.path.join(self.path, ".lock"), os.O_RDWR | os.O_CREAT, 0o600) + fcntl.flock(self._lock_fd, fcntl.LOCK_EX) + return self + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + assert self._lock_fd is not None + os.close(self._lock_fd) # closing releases the lock + self._lock_fd = None + + def _assert_locked(self) -> None: + assert self._lock_fd is not None, "PoolDirectory must be used as a context manager" + + # Path builders; these do not touch the filesystem and need no lock. + + def pending_path(self, uid: str) -> str: + return os.path.join(self.path, f"pending-{uid}.json") + + def conf_path(self, uid: str) -> str: + return os.path.join(self.path, f"conf-{uid}.json") + + def server_path(self, uid: str) -> str: + return os.path.join(self.path, f"server-{uid}.json") + + def claimed_path(self, client_pid: int, uid: str) -> str: + return os.path.join(self.path, f"claimed-{client_pid}-{uid}.json") + + def retired_path(self, uid: str) -> str: + return os.path.join(self.path, f"retired-{uid}.json") + + def member_dir(self, uid: str) -> str: + return os.path.join(self.path, f"member-{uid}") + + @staticmethod + def parse_entry(name: str) -> Tuple[Optional[str], Optional[str]]: + """The ``(kind, uid)`` of a pool directory entry, ``(None, None)`` for anything + else (the lock file, editor droppings, ...).""" + if name.startswith("member-"): + return "member", name[len("member-") :] + if not name.endswith(".json"): + return None, None + stem = name[: -len(".json")] + for kind in ("pending", "conf", "server", "retired"): + if stem.startswith(kind + "-"): + return kind, stem[len(kind) + 1 :] + if stem.startswith("claimed-"): + client_pid, sep, uid = stem[len("claimed-") :].partition("-") + if sep and client_pid.isdigit(): + return "claimed", uid + return None, None + + @staticmethod + def claiming_pid(claimed_path: str) -> int: + """The client pid recorded in a ``claimed--.json`` file name.""" + return int(os.path.basename(claimed_path)[len("claimed-") :].split("-", 1)[0]) + + # Locked accessors. + + def uids(self) -> List[str]: + self._assert_locked() + seen = [] + for name in self._entries(): + _, uid = self.parse_entry(name) + if uid is not None and uid not in seen: + seen.append(uid) + return seen + + def states(self, uid: str) -> Dict[str, str]: + """The state entries currently existing for ``uid``, as ``{kind: path}`` with kinds + ``pending``, ``conf``, ``server``, ``claimed``, ``retired``, and ``member`` (the + member's directory).""" + self._assert_locked() + found: Dict[str, str] = {} + for name in self._entries(): + kind, entry_uid = self.parse_entry(name) + if kind is not None and entry_uid == uid: + found[kind] = os.path.join(self.path, name) + return found + + def paths_of_kind(self, kind: str) -> List[Tuple[str, str]]: + """All ``(uid, path)`` of one state kind.""" + self._assert_locked() + return [ + (uid, os.path.join(self.path, name)) + for name in self._entries() + for entry_kind, uid in (self.parse_entry(name),) + if entry_kind == kind and uid is not None + ] + + def _entries(self) -> List[str]: + try: + return sorted(os.listdir(self.path)) + except OSError: + return [] + + def read_json(self, path: str) -> Optional[Dict[str, Any]]: + """``None`` for files that are missing or unreadable -- callers treat both like the + state not existing, and the reaping rules remove unreadable leftovers.""" + self._assert_locked() + try: + with open(path, "r") as f: + data = json.load(f) + except (OSError, ValueError): + return None + return data if isinstance(data, dict) else None + + def write_json(self, path: str, data: Dict[str, Any]) -> None: + self._assert_locked() + # 0600 like the reuse discovery file: server entries hold the auth token. + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as f: + os.fchmod(fd, 0o600) + f.write(json.dumps(data)) + + def rename(self, src: str, dst: str) -> None: + self._assert_locked() + os.rename(src, dst) + + def remove(self, path: str) -> None: + self._assert_locked() + with contextlib.suppress(FileNotFoundError): + os.remove(path) + + def remove_member_dir(self, uid: str) -> None: + self._assert_locked() + shutil.rmtree(self.member_dir(uid), ignore_errors=True) + + +class ServerPool: + """Claiming, refilling, and reaping the members of one pool directory.""" + + def __init__(self, directory: Optional[PoolDirectory] = None): + self._directory = directory or PoolDirectory() + + def acquire(self, master: str, opts: Dict[str, Any]) -> PoolMember: + """Claim one ready, fingerprint-matching member and top the pool back up. + + On a warm pool this returns after one janitor-claim-refill pass; on a cold pool + (first run, conf change, or all members consumed) the refill starts a full complement + and the loop waits for the first member to become ready, which costs one ordinary + cold start. The lock is released between polls so attendants can publish; launches + that die are relaunched by later passes, and only the overall deadline fails. + """ + from pyspark.sql.connect.local_server import startup_seed_conf + + seed_conf = startup_seed_conf(opts) + fingerprint = pool_fingerprint(master, seed_conf) + target = _pool_size(opts) + deadline = time.monotonic() + _ACQUIRE_TIMEOUT + while True: + with self._directory: + self.janitor() + member = self.claim(fingerprint) + self.refill(master, seed_conf, fingerprint, target) + if member is not None: + return member + if time.monotonic() >= deadline: + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={ + "reason": f"no pooled local server became ready within " + f"{_ACQUIRE_TIMEOUT}s; see the attendant and server logs under " + f"{self._directory.path}" + }, + ) + time.sleep(0.25) + + def claim(self, fingerprint: str) -> Optional[PoolMember]: + """Claim the oldest usable member with this fingerprint, or ``None``. The rename to + ``claimed--.json`` marks the member as owned by this process; the reaping + rules use that pid to retire members whose client died without releasing them.""" + candidates = [] + for uid, path in self._directory.paths_of_kind("server"): + data = self._directory.read_json(path) + if data is not None and data.get("fingerprint") == fingerprint: + candidates.append((PoolMember(data), uid, path)) + candidates.sort(key=lambda c: c[0].created) + for member, uid, path in candidates: + if not member.is_usable(): + continue # left for the reaping rules to retire + claim_path = self._directory.claimed_path(os.getpid(), uid) + self._directory.rename(path, claim_path) + member.claim_path = claim_path + return member + return None + + def refill(self, master: str, seed_conf: Dict[str, Any], fingerprint: str, target: int) -> None: + """Launch members until ready or in-flight ones with this fingerprint reach + ``target``. Running under the directory lock is what makes concurrent cold starters + share one complement of launches instead of each spawning their own.""" + available = 0 + for kind in ("server", "pending"): + for _, path in self._directory.paths_of_kind(kind): + data = self._directory.read_json(path) + if data is not None and data.get("fingerprint") == fingerprint: + available += 1 + for _ in range(target - available): + MemberAttendant.spawn(self._directory, master, seed_conf, fingerprint) + + def janitor(self) -> None: + """Reap leftovers of launches, clients, and attendants that died uncleanly. Every + rule is idempotent, so successive passes from any process are safe.""" + for uid in self._directory.uids(): + self.reap(uid) + + def reap(self, uid: str) -> bool: + """Apply the reaping rules to one member; ``True`` when nothing of it remains. + Shared by the janitor (all members) and by each attendant supervising its own member. + """ + states = self._directory.states(uid) + if "pending" in states: + self._reap_pending(uid, states["pending"]) + if "server" in states: + self._reap_server(states["server"]) + if "claimed" in states: + self._reap_claimed(states["claimed"]) + if "retired" in states: + self._reap_retired(uid, states["retired"]) + + remaining = self._directory.states(uid) + if set(remaining) == {"member"}: + # Nothing references the member directory anymore. The age gate keeps the logs + # of a freshly failed launch around long enough to be looked at. + try: + expired = time.time() - os.path.getmtime(remaining["member"]) > _DIR_GC_AGE + except OSError: + expired = True + if expired: + self._directory.remove_member_dir(uid) + remaining = self._directory.states(uid) + return not remaining + + def _reap_pending(self, uid: str, path: str) -> None: + """A launch whose attendant died or hung: kill the attendant and whatever server + spark-daemon.sh may have recorded for it, and withdraw the launch's bookkeeping so + refills stop counting it.""" + data = self._directory.read_json(path) + attendant_pid = int(data["attendant_pid"]) if data else -1 + age = time.time() - float(data["created"]) if data else _LAUNCH_TIMEOUT + 1 + if not _pid_alive(attendant_pid) or age > _LAUNCH_TIMEOUT: + _signal(attendant_pid, signal.SIGTERM) + self.abort_launch(uid) + + def abort_launch(self, uid: str) -> None: + """Withdraw a failed launch and stop any server it started before failing.""" + self._kill_recorded_daemon(uid) + self._directory.remove(self._directory.pending_path(uid)) + self._directory.remove(self._directory.conf_path(uid)) + + def _reap_server(self, path: str) -> None: + """A ready member that is unusable (dead, unreachable, version-mismatched after an + upgrade, or an unreadable record) or has sat unclaimed past the idle timeout: retire + it.""" + data = self._directory.read_json(path) + member = PoolMember(data) if data is not None else None + idle = _idle_timeout() + expired = member is not None and idle > 0 and time.time() - member.created > idle + if member is None or expired or not member.is_usable(): + self._retire(path, member.pid if member is not None else -1) + + def _reap_claimed(self, path: str) -> None: + """A claimed member whose client died without releasing it (e.g. SIGKILL), or whose + server died under its client: retire it. Claims of this live process are its own.""" + data = self._directory.read_json(path) + server_pid = int(data["pid"]) if data else -1 + client_pid = self._directory.claiming_pid(path) + if client_pid == os.getpid(): + return + if not _pid_alive(client_pid) or not _pid_alive(server_pid): + self._retire(path, server_pid) + + def _reap_retired(self, uid: str, path: str) -> None: + """A retiring member: drop it once its server is gone, hard-kill the server if it + hangs in shutdown, and eventually stop tracking one that survives even that (it is + not ours to signal; nothing more can be done).""" + data = self._directory.read_json(path) + server_pid = int(data["pid"]) if data else -1 + age = time.time() - float(data["retired"]) if data else _RETIRE_GIVE_UP + 1 + if not _pid_alive(server_pid) or age > _RETIRE_GIVE_UP: + self._directory.remove(path) + self._directory.remove_member_dir(uid) + elif age > _RETIRE_KILL_AFTER: + _signal(server_pid, signal.SIGKILL) + + def _retire(self, state_path: str, server_pid: int) -> None: + """Move a member into the retired state: signal its server and track the shutdown so + :meth:`_reap_retired` can escalate if the JVM hangs.""" + _, uid = self._directory.parse_entry(os.path.basename(state_path)) + assert uid is not None + _signal(server_pid, signal.SIGTERM) + self._directory.remove(state_path) + self._directory.write_json( + self._directory.retired_path(uid), {"pid": server_pid, "retired": time.time()} + ) + + def _kill_recorded_daemon(self, uid: str) -> None: + """Kill the server pid that spark-daemon.sh recorded in the member directory, if + any. This is how a half-started JVM is reaped when its launch fails or its attendant + dies before publishing.""" + from pyspark.sql.connect.local_server import Discovery + + discovery = Discovery(os.path.join(self._directory.member_dir(uid), "connect-local.json")) + daemon_pid = discovery.daemon_pid() + if daemon_pid is not None: + _signal(daemon_pid, signal.SIGTERM) + + def release(self, member: PoolMember) -> None: + """Retire this process's claimed member; the shutdown completes in the background, + watched by the member's attendant with the janitor as backstop.""" + assert member.claim_path is not None + with self._directory: + self._retire(member.claim_path, member.pid) + + def purge(self) -> int: + """Force-stop every member -- ready, in-flight, or claimed -- and empty the pool + directory; the escape hatch back to a clean slate. Returns the number of processes + signalled. SIGKILL rather than SIGTERM because nothing tracks a member once its + state files are gone, so a shutdown that hangs would leak. Supervising attendants + hold no state file; they notice the emptied directory and exit on their own.""" + signalled = 0 + with self._directory: + for uid in self._directory.uids(): + for kind, path in self._directory.states(uid).items(): + if kind == "member": + continue + data = self._directory.read_json(path) + for key in ("pid", "attendant_pid"): + if data is not None and key in data: + if _signal(int(data[key]), signal.SIGKILL): + signalled += 1 + self._directory.remove(path) + self._kill_recorded_daemon(uid) + self._directory.remove_member_dir(uid) + return signalled + + +# The member this client process has claimed, if any. Module-level so the session's stop +# callback and atexit share one idempotent release path. +_claimed_member: Optional[PoolMember] = None +_release_registered = False + + +def acquire_pooled_local_connect_server(master: str, opts: Dict[str, Any]) -> str: + """Claim a local Connect server not previously assigned to an application run. + + Returns the ``sc://host:port`` endpoint and sets ``SPARK_CONNECT_AUTHENTICATE_TOKEN`` so + the client authenticates against that server. Only reached for a ``local`` master when the + pool opt-in is set; see ``SparkSession.getOrCreate`` in ``pyspark.sql.session``. + """ + global _claimed_member, _release_registered + if os.name != "posix": + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={ + "reason": "spark.local.connect.pool relies on the POSIX scripts under sbin/; " + "on this platform start a server manually (sbin/start-connect-server.sh) and " + 'connect with .remote("sc://...")' + }, + ) + # getOrCreate() may be re-entered while this process already holds a live claimed member + # (the connect layer then returns the existing session); claiming again would strand a + # second server. + if _claimed_member is not None and _pid_alive(_claimed_member.pid): + os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = _claimed_member.token + return _claimed_member.url + + member = ServerPool().acquire(master, opts) + _claimed_member = member + if not _release_registered: + # A client that exits without stopping its session still releases its member; the + # member's attendant and the janitor cover clients that die uncleanly. + atexit.register(release_pooled_local_connect_server) + _release_registered = True + os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = member.token + return member.url + + +def release_pooled_local_connect_server() -> None: + """Retire this process's claimed pooled server; safe to call when there is none. The + server winds down in the background while this client moves on.""" + global _claimed_member + member, _claimed_member = _claimed_member, None + if member is not None: + ServerPool().release(member) + + +def purge_local_connect_pool() -> int: + """See :meth:`ServerPool.purge`. Also available as + ``python -m pyspark.sql.connect.local_server_pool --purge``.""" + return ServerPool().purge() + + +class MemberAttendant: + """Boots one pool member, publishes it, warms it, and supervises it until it is gone. + + Runs detached from the spawning client (``--attend``). Every phase is appended to + ``member-/attendant.log`` so a member that misbehaved can be debugged after the + fact. A failed boot reaps whatever half-started server spark-daemon.sh recorded and + withdraws the launch's bookkeeping, so refills stop counting it. + """ + + @classmethod + def spawn( + cls, + directory: PoolDirectory, + master: str, + seed_conf: Dict[str, Any], + fingerprint: str, + ) -> None: + """Start one detached attendant and publish its in-flight launch. + + Callers hold ``directory``'s lock, so another refill sees the pending marker before it + can start a duplicate launch for the same pool complement. + """ + uid = uuid.uuid4().hex[:12] + directory.write_json(directory.conf_path(uid), seed_conf) + cmd = [ + sys.executable, + "-m", + "pyspark.sql.connect.local_server_pool", + "--attend", + "--pool-dir", + directory.path, + "--uid", + uid, + "--master", + master, + "--fingerprint", + fingerprint, + ] + env = dict(os.environ) + # The attendant must neither see this client's Connect mode nor inherit its auth + # token: each member gets its own token from LocalConnectServer. + for var in ( + "SPARK_REMOTE", + "SPARK_LOCAL_REMOTE", + "SPARK_CONNECT_MODE_ENABLED", + "SPARK_CONNECT_AUTHENTICATE_TOKEN", + ): + env.pop(var, None) + try: + proc = subprocess.Popen( + cmd, + env=env, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + except OSError: + directory.remove(directory.conf_path(uid)) + raise + directory.write_json( + directory.pending_path(uid), + {"attendant_pid": proc.pid, "created": time.time(), "fingerprint": fingerprint}, + ) + + def __init__(self, directory: PoolDirectory, uid: str, master: str, fingerprint: str): + self._directory = directory + self._pool = ServerPool(directory) + self._uid = uid + self._master = master + self._fingerprint = fingerprint + self._member_dir = directory.member_dir(uid) + + def run(self) -> int: + os.makedirs(self._member_dir, mode=0o700, exist_ok=True) + member = self._boot() + if member is None: + return 1 + if _warmup_enabled(): + self._warm(member) + self._log("supervising") + self._supervise(member.pid) + self._log("done") + return 0 + + def _log(self, message: str) -> None: + with open(os.path.join(self._member_dir, "attendant.log"), "a") as f: + f.write(f"{time.time():.3f} {message}\n") + + def _boot(self) -> Optional[PoolMember]: + """Start the server on a fresh ephemeral port and publish it as ready-to-claim. + Publishing consumes the launch's pending marker and conf seed: from that moment the + member counts as a server, and this process's job shifts to supervising it.""" + from pyspark.sql.connect.local_server import Discovery, LocalConnectServer + + with self._directory: + seed_conf = self._directory.read_json(self._directory.conf_path(self._uid)) or {} + discovery = Discovery(os.path.join(self._member_dir, "connect-local.json")) + self._log(f"booting master={self._master}") + try: + with discovery: + server = LocalConnectServer(discovery) + server.start( + self._master, + {}, + use_ephemeral_port=True, + seed_conf=seed_conf, + ) + data = discovery.load() + assert data is not None # launch() saved it + except Exception as e: + self._log(f"boot failed: {e!r}") + with self._directory: + self._pool.abort_launch(self._uid) + return None + data["fingerprint"] = self._fingerprint + data["created"] = time.time() + with self._directory: + self._directory.write_json(self._directory.server_path(self._uid), data) + self._directory.remove(self._directory.pending_path(self._uid)) + self._directory.remove(self._directory.conf_path(self._uid)) + self._log(f"published pid={data['pid']} port={data['port']}") + return PoolMember(data) + + def _supervise(self, server_pid: int) -> None: + """Watch this member until nothing of it remains, applying the pool's reaping rules. + + This is what lets an idle machine drain to zero servers with no further Spark run: the + idle timeout, dead-client cleanup, and hard-kill escalation all fire from here even + when no client ever runs again. + """ + while True: + with self._directory: + if self._pool.reap(self._uid): + return + states = set(self._directory.states(self._uid)) + if states <= {"member"} and not _pid_alive(server_pid): + # A purge or another process's janitor already tore the member down; only + # the young member directory remains. + self._directory.remove_member_dir(self._uid) + return + time.sleep(2) + + def _warm(self, member: PoolMember) -> None: + """Run the fixed warmup queries in a bounded child process. + + A child rather than a thread so that a member claimed, used, and torn down mid-warmup + cannot wedge this attendant: a gRPC call against a dying JVM can block indefinitely, + and a stuck child is simply killed. The member is published before warmup starts, so + a client in a hurry can claim it mid-warmup; the child is stopped the moment that + happens. + """ + env = dict(os.environ) + env["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = member.token + for var in ("SPARK_REMOTE", "SPARK_LOCAL_REMOTE", "SPARK_CONNECT_MODE_ENABLED"): + env.pop(var, None) + self._log("warmup starting") + child = subprocess.Popen( + [ + sys.executable, + "-m", + "pyspark.sql.connect.local_server_pool", + "--warm", + "--url", + member.url, + ], + env=env, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + deadline = time.monotonic() + 120 + while child.poll() is None: + # A plain existence poll, not a locked read: the only cost of a stale result is + # stopping the child one tick later. + claimed = not os.path.exists(self._directory.server_path(self._uid)) + if claimed or time.monotonic() > deadline: + child.kill() + child.wait() + self._log("warmup stopped (member claimed or warmup timed out)") + return + time.sleep(0.5) + self._log(f"warmup finished with code {child.returncode}") + + +def _run_warmup(url: str) -> None: + """The fixed synthetic warmup, run inside the ``--warm`` child against one member. + + JIT and codegen caches are JVM-global, so warming them here removes most of the + first-query latency from the one real run this member will serve. The queries create no + catalog objects, temp views, or cached data, and the warmup session is released at the + end, so nothing of it outlives this process. + """ + from pyspark.sql.connect.session import SparkSession as ConnectSession + + session = ConnectSession(connection=url) + try: + session.sql("SELECT 1").collect() + session.sql("SELECT sum(id) AS s, count(1) AS c FROM range(100000)").collect() + session.sql( + "SELECT x, count(*) AS c FROM (SELECT id % 7 AS x FROM range(100000)) GROUP BY x" + ).collect() + session.range(100000).selectExpr("sum(id)").collect() + finally: + session.stop() + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Manage the opt-in pool of pre-warmed single-use local Spark Connect " + "servers (spark.local.connect.pool)." + ) + parser.add_argument( + "--purge", + action="store_true", + help="force-stop every pool member and empty the pool directory", + ) + # Internal entry points, spawned by ServerPool and MemberAttendant respectively. + parser.add_argument("--attend", action="store_true", help=argparse.SUPPRESS) + parser.add_argument("--pool-dir", help=argparse.SUPPRESS) + parser.add_argument("--uid", help=argparse.SUPPRESS) + parser.add_argument("--master", help=argparse.SUPPRESS) + parser.add_argument("--fingerprint", help=argparse.SUPPRESS) + parser.add_argument("--warm", action="store_true", help=argparse.SUPPRESS) + parser.add_argument("--url", help=argparse.SUPPRESS) + args = parser.parse_args() + + if args.purge: + print(f"Signalled {purge_local_connect_pool()} pool process(es).") + elif args.attend: + attendant = MemberAttendant( + PoolDirectory(args.pool_dir), args.uid, args.master, args.fingerprint + ) + sys.exit(attendant.run()) + elif args.warm: + _run_warmup(args.url) + else: + parser.print_help(sys.stderr) + sys.exit(2) + + +if __name__ == "__main__": + main() diff --git a/python/pyspark/sql/connect/session.py b/python/pyspark/sql/connect/session.py index 1a3e1c1aebb39..8e1f7e46f30ce 100644 --- a/python/pyspark/sql/connect/session.py +++ b/python/pyspark/sql/connect/session.py @@ -310,8 +310,19 @@ def __init__( ) self._session_id = self._client._session_id + self._initialize_lifecycle_state() + + def _initialize_lifecycle_state(self) -> None: + """Initialize state shared by normally constructed and derived sessions.""" + # Set to false to prevent client.release_session on close() (testing only) self.release_session_on_close = True + self._on_stop_callbacks: List[Callable[[], None]] = [] + + def _register_on_stop_callback(self, callback: Callable[[], None]) -> None: + """Register internal lifecycle cleanup owned by the code that created this session.""" + if callback not in self._on_stop_callbacks: + self._on_stop_callbacks.append(callback) @classmethod def _set_default_and_active_session(cls, session: "SparkSession") -> None: @@ -991,6 +1002,13 @@ def stop(self) -> None: if "SPARK_REMOTE" in os.environ: del os.environ["SPARK_REMOTE"] + callbacks, self._on_stop_callbacks = self._on_stop_callbacks, [] + for callback in callbacks: + try: + callback() + except Exception as e: + logger.warning(f"session.stop(): Cleanup callback failed. Error: {e}") + def __enter__(self) -> "SparkSession": """ Enable 'with SparkSession.builder.(...).getOrCreate() as session: app' syntax. @@ -1370,7 +1388,7 @@ def cloneSession(self, new_session_id: Optional[str] = None) -> "SparkSession": new_session = object.__new__(SparkSession) new_session._client = cloned_client new_session._session_id = cloned_client._session_id - new_session.release_session_on_close = True + new_session._initialize_lifecycle_state() return new_session def newSession(self) -> "SparkSession": @@ -1398,7 +1416,7 @@ def newSession(self) -> "SparkSession": new_session = object.__new__(SparkSession) new_session._client = new_client new_session._session_id = new_client._session_id - new_session.release_session_on_close = True + new_session._initialize_lifecycle_state() return new_session diff --git a/python/pyspark/sql/session.py b/python/pyspark/sql/session.py index e9398bd512782..253d1a83508e6 100644 --- a/python/pyspark/sql/session.py +++ b/python/pyspark/sql/session.py @@ -523,6 +523,14 @@ def getOrCreate(self) -> "SparkSession": messageParameters={}, ) + pool_cleanup: Optional[Callable[[], None]] = None + pool_local = str( + opts.get( + "spark.local.connect.pool", + os.environ.get("SPARK_LOCAL_CONNECT_POOL", ""), + ) + ).lower() in ("1", "true") + reuse_local = str( opts.get( "spark.local.connect.reuse", @@ -530,7 +538,23 @@ def getOrCreate(self) -> "SparkSession": ) ).lower() in ("1", "true") - if url.startswith("local") and reuse_local: + if url.startswith("local") and pool_local: + from pyspark.sql.connect.local_server_pool import ( + acquire_pooled_local_connect_server, + release_pooled_local_connect_server, + ) + + # Opt-in: claim a server not previously assigned to an + # application run from a pool of pre-warmed local Connect + # servers. It is torn down when this run's session stops, so no + # state carries across runs. Takes precedence over reuse. See + # `pyspark.sql.connect.local_server_pool`. + url = acquire_pooled_local_connect_server(url, opts) + pool_cleanup = release_pooled_local_connect_server + for k in list(opts): + if k.startswith("spark.local.connect."): + opts.pop(k) + elif url.startswith("local") and reuse_local: from pyspark.sql.connect.local_server import ( reuse_or_start_local_connect_server, ) @@ -551,9 +575,27 @@ def getOrCreate(self) -> "SparkSession": os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1" opts["spark.remote"] = url + try: + remote_session = RemoteSparkSession.builder.config( + map=opts + ).getOrCreate() + except Exception: + if pool_cleanup is not None: + try: + pool_cleanup() + except Exception as cleanup_error: + warnings.warn( + "Failed to clean up a pooled local Connect server " + f"after session creation failed: {cleanup_error}", + RuntimeWarning, + stacklevel=2, + ) + raise + if pool_cleanup is not None: + remote_session._register_on_stop_callback(pool_cleanup) return cast( SparkSession, - RemoteSparkSession.builder.config(map=opts).getOrCreate(), + remote_session, ) elif "SPARK_LOCAL_REMOTE" in os.environ: url = "sc://localhost" 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 eba26972ae8c3..857ea02e4060e 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server.py @@ -114,6 +114,65 @@ def test_discovery_location(self) -> None: self.assertIn("spark-connect-{}".format(getpass.getuser()), default.directory) self.assertEqual(os.stat(default.directory).st_mode & 0o777, 0o700) + def test_startup_seed_conf(self) -> None: + from unittest import mock + + initial = { + "spark.sql.shuffle.partitions": "8", + "spark.master": "local[1]", + } + opts = { + "spark.sql.warehouse.dir": os.path.join(self._tmpdir, "warehouse"), + "spark.local.connect.reuse": "true", + "spark.connect.grpc.binding.port": "0", + } + env = { + "PYSPARK_REMOTE_INIT_CONF_LEN": "1", + "PYSPARK_REMOTE_INIT_CONF_0": json.dumps(initial), + } + with mock.patch.dict(os.environ, env): + self.assertEqual( + local_server.startup_seed_conf(opts), + { + "spark.sql.shuffle.partitions": "8", + "spark.sql.warehouse.dir": opts["spark.sql.warehouse.dir"], + }, + ) + + def test_start_delegates_launch_options(self) -> None: + from unittest import mock + + discovery = mock.Mock() + discovery.load.side_effect = [ + None, + { + "host": "localhost", + "port": 15002, + "token": "t", + "pid": os.getpid(), + "spark_version": __version__, + }, + ] + server = LocalConnectServer(discovery) + seed_conf = {"spark.sql.shuffle.partitions": "4"} + with mock.patch.object(local_server, "ServerLauncher") as launcher: + server.start( + "local[2]", + {"spark.local.connect.reuse": "true"}, + use_ephemeral_port=True, + seed_conf=seed_conf, + ) + + launcher.assert_called_once_with( + "local[2]", + {"spark.local.connect.reuse": "true"}, + discovery, + use_ephemeral_port=True, + seed_conf=seed_conf, + ) + launcher.return_value.launch.assert_called_once_with() + self.assertEqual(server.port, 15002) + def test_discovery_roundtrip(self) -> None: with Discovery() as discovery: saved = self._server(port=15002) diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py new file mode 100644 index 0000000000000..a3945cac5886e --- /dev/null +++ b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py @@ -0,0 +1,645 @@ +# +# 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 contextlib +import json +import os +import shutil +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.connect import local_server_pool + from pyspark.sql.connect.local_server_pool import ( + MemberAttendant, + PoolDirectory, + PoolMember, + ServerPool, + pool_fingerprint, + ) + from pyspark.version import __version__ + + +@contextlib.contextmanager +def _listening_socket(): + import socket + + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + listener.bind(("localhost", 0)) + listener.listen(1) + yield listener.getsockname()[1] + finally: + listener.close() + + +def _closed_port() -> int: + """A port with nothing listening on it.""" + import socket + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("localhost", 0)) + return sock.getsockname()[1] + + +def _spawn_sleeper() -> "subprocess.Popen": + """A long sleeper standing in for a pool server or attendant process.""" + return subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(300)"], + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + + +def _spawn_stubborn_sleeper() -> "subprocess.Popen": + """A sleeper that ignores SIGTERM, standing in for a server hanging in shutdown. It + prints one line once its handler is installed so tests do not signal it too early.""" + proc = subprocess.Popen( + [ + sys.executable, + "-c", + "import signal, time\n" + "signal.signal(signal.SIGTERM, signal.SIG_IGN)\n" + "print('ready', flush=True)\n" + "time.sleep(300)", + ], + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + ) + assert proc.stdout is not None + proc.stdout.readline() + return proc + + +def _wait_proc_dead(proc: "subprocess.Popen", timeout: float = 30.0) -> bool: + try: + proc.wait(timeout=timeout) + return True + except subprocess.TimeoutExpired: + return False + + +def _wait_pid_gone(pid: int, timeout: float = 60.0) -> bool: + """Wait for a non-child pid to stop running. Linux zombies count as terminated.""" + deadline = time.time() + timeout + while time.time() < deadline: + if not local_server_pool._pid_alive(pid): + return True + time.sleep(0.2) + return False + + +_SAVED_ENV_KEYS = ( + "SPARK_LOCAL_CONNECT_POOL", + "SPARK_LOCAL_CONNECT_POOL_DIR", + "SPARK_LOCAL_CONNECT_POOL_SIZE", + "SPARK_LOCAL_CONNECT_POOL_IDLE_TIMEOUT", + "SPARK_LOCAL_CONNECT_POOL_WARMUP", + "SPARK_CONNECT_AUTHENTICATE_TOKEN", + "PYSPARK_PYTHON", +) + + +@unittest.skipIf( + not should_test_connect or is_remote_only(), + connect_requirement_message or "Requires Spark Connect test dependencies", +) +class LocalConnectServerPoolUnitTests(unittest.TestCase): + """Tests for the pool filesystem model; no real servers are started.""" + + def setUp(self) -> None: + self._tmpdir = tempfile.mkdtemp() + self._saved_env = {k: os.environ.get(k) for k in _SAVED_ENV_KEYS} + for k in _SAVED_ENV_KEYS: + os.environ.pop(k, None) + os.environ["SPARK_LOCAL_CONNECT_POOL_DIR"] = os.path.join(self._tmpdir, "pool") + self._directory = PoolDirectory() + self._pool = ServerPool(self._directory) + self._procs = [] + local_server_pool._claimed_member = None + + def tearDown(self) -> None: + local_server_pool._claimed_member = None + for proc in self._procs: + try: + proc.kill() + proc.communicate(timeout=10) + except Exception: + pass + for k, v in self._saved_env.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + shutil.rmtree(self._tmpdir, ignore_errors=True) + + def _sleeper(self) -> "subprocess.Popen": + proc = _spawn_sleeper() + self._procs.append(proc) + return proc + + def _stubborn_sleeper(self) -> "subprocess.Popen": + proc = _spawn_stubborn_sleeper() + self._procs.append(proc) + return proc + + def _server_data(self, port: int, pid: int, fingerprint: str = "fp", **overrides) -> dict: + data = { + "host": "localhost", + "port": port, + "token": "t", + "pid": pid, + "spark_version": __version__, + "fingerprint": fingerprint, + "created": time.time(), + } + data.update(overrides) + return data + + def _write_state(self, path: str, data: dict) -> str: + with self._directory as directory: + directory.write_json(path, data) + return path + + def _states(self, uid: str) -> dict: + with self._directory as directory: + return directory.states(uid) + + def test_pool_directory_location(self) -> None: + self.assertEqual(self._directory.path, os.path.join(self._tmpdir, "pool")) + os.environ.pop("SPARK_LOCAL_CONNECT_POOL_DIR") + default = PoolDirectory() + self.assertEqual(os.path.basename(default.path), "pool") + self.assertTrue(default.path.startswith(tempfile.gettempdir())) + + def test_pool_size_parsing(self) -> None: + from pyspark.sql.connect.local_server_pool import _pool_size + + self.assertEqual(_pool_size({}), 2) + self.assertEqual(_pool_size({"spark.local.connect.pool.size": "3"}), 3) + os.environ["SPARK_LOCAL_CONNECT_POOL_SIZE"] = "5" + self.assertEqual(_pool_size({}), 5) + # Junk falls back to the default; values below one are clamped up. + self.assertEqual(_pool_size({"spark.local.connect.pool.size": "abc"}), 2) + self.assertEqual(_pool_size({"spark.local.connect.pool.size": "0"}), 1) + + def test_warmup_enabled_parsing(self) -> None: + from pyspark.sql.connect.local_server_pool import _warmup_enabled + + self.assertTrue(_warmup_enabled()) + for value in ("0", "false", "FALSE"): + os.environ["SPARK_LOCAL_CONNECT_POOL_WARMUP"] = value + self.assertFalse(_warmup_enabled()) + os.environ["SPARK_LOCAL_CONNECT_POOL_WARMUP"] = "1" + self.assertTrue(_warmup_enabled()) + + @unittest.skipUnless( + sys.platform.startswith("linux") and os.path.isdir("/proc"), + "requires Linux process state", + ) + def test_pid_alive_treats_zombie_as_dead(self) -> None: + proc = subprocess.Popen([sys.executable, "-c", "pass"]) + try: + state = "" + deadline = time.time() + 5 + while time.time() < deadline: + with open(f"/proc/{proc.pid}/stat", encoding="utf-8") as stat_file: + fields = stat_file.read().rpartition(")")[2].split() + state = fields[0] if fields else "" + if state == "Z": + break + time.sleep(0.01) + self.assertEqual(state, "Z", "the child did not enter zombie state") + self.assertFalse(local_server_pool._pid_alive(proc.pid)) + finally: + proc.wait(timeout=10) + + def test_fingerprint_identity(self) -> None: + base = pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"}) + self.assertEqual(base, pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"})) + self.assertNotEqual( + base, pool_fingerprint("local[2]", {"spark.sql.shuffle.partitions": "4"}) + ) + self.assertNotEqual( + base, pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "8"}) + ) + self.assertNotEqual(base, pool_fingerprint("local[*]", {})) + # The working directory shapes the server (relative warehouse and metastore paths), + # so members are never shared across directories. + cwd = os.getcwd() + try: + os.chdir(self._tmpdir) + self.assertNotEqual(base, pool_fingerprint("local[*]", {"x": "4"})) + in_tmpdir = pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"}) + finally: + os.chdir(cwd) + self.assertNotEqual(base, in_tmpdir) + # So does the Python environment the server would run UDFs with. + os.environ["PYSPARK_PYTHON"] = "/some/other/python" + self.assertNotEqual( + base, pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"}) + ) + + def test_claim_matches_fingerprint_and_renames(self) -> None: + with _listening_socket() as port: + sleeper = self._sleeper() + self._write_state( + self._directory.server_path("aaa"), + self._server_data(port, sleeper.pid, fingerprint="other-fp"), + ) + self._write_state( + self._directory.server_path("bbb"), + self._server_data(port, sleeper.pid, fingerprint="my-fp", token="t-bbb"), + ) + with self._directory: + member = self._pool.claim("my-fp") + self.assertIsNotNone(member) + self.assertEqual(member.token, "t-bbb") + claim_name = f"claimed-{os.getpid()}-bbb.json" + self.assertEqual(os.path.basename(member.claim_path), claim_name) + states = self._states("bbb") + self.assertEqual(set(states), {"claimed"}) + # The mismatched member is untouched, and a second claim finds nothing. + self.assertEqual(set(self._states("aaa")), {"server"}) + with self._directory: + self.assertIsNone(self._pool.claim("my-fp")) + + def test_claim_prefers_the_oldest_member(self) -> None: + with _listening_socket() as port: + sleeper = self._sleeper() + for uid, created in (("young", time.time()), ("old", time.time() - 100)): + self._write_state( + self._directory.server_path(uid), + self._server_data(port, sleeper.pid, token="t-" + uid, created=created), + ) + with self._directory: + member = self._pool.claim("fp") + # The oldest member has had the longest to finish its warmup. + self.assertEqual(member.token, "t-old") + + def test_claim_skips_unreachable_member(self) -> None: + self._write_state( + self._directory.server_path("ccc"), self._server_data(_closed_port(), os.getpid()) + ) + with self._directory: + self.assertIsNone(self._pool.claim("fp")) + + def test_reap_pending_of_dead_attendant(self) -> None: + # The attendant died mid-boot: its pending marker and conf seed are withdrawn, and + # the half-started server whose pid spark-daemon.sh recorded is killed. + from pyspark.sql.connect.local_server import Discovery + + half_started = self._sleeper() + member_dir = self._directory.member_dir("boot") + os.makedirs(member_dir) + discovery = Discovery(os.path.join(member_dir, "connect-local.json")) + with open(discovery.daemon_pid_path, "w") as f: + f.write(str(half_started.pid)) + self._write_state( + self._directory.pending_path("boot"), + {"attendant_pid": 2**31 - 1, "created": time.time(), "fingerprint": "fp"}, + ) + self._write_state(self._directory.conf_path("boot"), {"spark.foo": "bar"}) + + with self._directory: + self._pool.reap("boot") + + states = self._states("boot") + self.assertNotIn("pending", states) + self.assertNotIn("conf", states) + self.assertTrue(_wait_proc_dead(half_started), "the half-started server was not reaped") + + def test_reap_keeps_live_pending(self) -> None: + attendant = self._sleeper() + self._write_state( + self._directory.pending_path("live"), + {"attendant_pid": attendant.pid, "created": time.time(), "fingerprint": "fp"}, + ) + with self._directory: + self._pool.reap("live") + self.assertIn("pending", self._states("live")) + self.assertIsNone(attendant.poll()) + + def test_reap_server_unreachable_and_idle(self) -> None: + with self.subTest("unreachable member is retired"): + gone = self._sleeper() + self._write_state( + self._directory.server_path("dead"), self._server_data(_closed_port(), gone.pid) + ) + with self._directory: + self._pool.reap("dead") + self.assertEqual(set(self._states("dead")), {"retired"}) + self.assertTrue(_wait_proc_dead(gone)) + with self.subTest("member idle past the timeout is retired"): + os.environ["SPARK_LOCAL_CONNECT_POOL_IDLE_TIMEOUT"] = "10" + with _listening_socket() as port: + idle = self._sleeper() + self._write_state( + self._directory.server_path("idle"), + self._server_data(port, idle.pid, created=time.time() - 60), + ) + fresh = self._sleeper() + self._write_state( + self._directory.server_path("fresh"), self._server_data(port, fresh.pid) + ) + with self._directory: + self._pool.janitor() + self.assertEqual(set(self._states("idle")), {"retired"}) + self.assertEqual(set(self._states("fresh")), {"server"}) + self.assertTrue(_wait_proc_dead(idle)) + self.assertIsNone(fresh.poll()) + + def test_reap_claimed_of_dead_client(self) -> None: + orphan = self._sleeper() + dead_client = 2**31 - 1 + self._write_state( + self._directory.claimed_path(dead_client, "orphan"), + self._server_data(_closed_port(), orphan.pid), + ) + ours = self._sleeper() + self._write_state( + self._directory.claimed_path(os.getpid(), "mine"), + self._server_data(_closed_port(), ours.pid), + ) + with self._directory: + self._pool.janitor() + # The orphaned claim is retired and its server stopped; our own claim is untouched. + self.assertEqual(set(self._states("orphan")), {"retired"}) + self.assertTrue(_wait_proc_dead(orphan), "the orphaned server was not stopped") + self.assertEqual(set(self._states("mine")), {"claimed"}) + self.assertIsNone(ours.poll()) + + def test_reap_retired_escalates_to_sigkill(self) -> None: + with self.subTest("a fresh retirement is left to shut down gracefully"): + fresh = self._sleeper() + self._write_state( + self._directory.retired_path("fresh"), + {"pid": fresh.pid, "retired": time.time()}, + ) + with self._directory: + self._pool.reap("fresh") + self.assertEqual(set(self._states("fresh")), {"retired"}) + self.assertIsNone(fresh.poll()) + with self.subTest("a hung shutdown is hard-killed"): + stubborn = self._stubborn_sleeper() + self._write_state( + self._directory.retired_path("hung"), + {"pid": stubborn.pid, "retired": time.time() - 31}, + ) + with self._directory: + self._pool.reap("hung") + self.assertTrue(_wait_proc_dead(stubborn), "SIGKILL escalation did not happen") + with self._directory: + self.assertTrue(self._pool.reap("hung")) + + def test_release_retires_the_claimed_member(self) -> None: + server = self._sleeper() + claim_path = self._write_state( + self._directory.claimed_path(os.getpid(), "xyz"), + self._server_data(_closed_port(), server.pid), + ) + member = PoolMember(self._server_data(_closed_port(), server.pid)) + member.claim_path = claim_path + local_server_pool._claimed_member = member + + local_server_pool.release_pooled_local_connect_server() + + self.assertIsNone(local_server_pool._claimed_member) + states = self._states("xyz") + self.assertEqual(set(states), {"retired"}) + with self._directory as directory: + self.assertEqual(directory.read_json(states["retired"])["pid"], server.pid) + self.assertTrue(_wait_proc_dead(server), "release did not stop the server") + # Releasing again is a no-op. + local_server_pool.release_pooled_local_connect_server() + + def test_purge_kills_everything_and_empties_the_directory(self) -> None: + warm = self._sleeper() + attendant = self._sleeper() + self._write_state( + self._directory.server_path("warm"), self._server_data(_closed_port(), warm.pid) + ) + self._write_state( + self._directory.pending_path("boot"), + {"attendant_pid": attendant.pid, "created": time.time(), "fingerprint": "fp"}, + ) + self._write_state(self._directory.conf_path("boot"), {"spark.foo": "bar"}) + os.makedirs(self._directory.member_dir("warm")) + + signalled = local_server_pool.purge_local_connect_pool() + + self.assertGreaterEqual(signalled, 2) + self.assertEqual(os.listdir(self._directory.path), [".lock"]) + self.assertTrue(_wait_proc_dead(warm)) + self.assertTrue(_wait_proc_dead(attendant)) + + def test_refill_only_counts_matching_members(self) -> None: + from unittest import mock + + with self._directory as directory: + directory.write_json( + directory.server_path("srv"), self._server_data(1, os.getpid(), "my-fp") + ) + directory.write_json( + directory.pending_path("pen"), + {"attendant_pid": os.getpid(), "created": time.time(), "fingerprint": "other"}, + ) + with mock.patch.object(MemberAttendant, "spawn") as spawn: + self._pool.refill("local[2]", {}, "my-fp", target=2) + # One matching member exists (the other-fingerprint launch does not count), so one + # launch tops the pool up to the target of two. + self.assertEqual(spawn.call_count, 1) + + def test_acquire_returns_the_member_already_claimed_by_this_process(self) -> None: + member = PoolMember(self._server_data(15002, os.getpid())) + member.claim_path = "unused" + local_server_pool._claimed_member = member + url = local_server_pool.acquire_pooled_local_connect_server("local[2]", {}) + self.assertEqual(url, "sc://localhost:15002") + self.assertEqual(os.environ.get("SPARK_CONNECT_AUTHENTICATE_TOKEN"), "t") + + def test_acquire_requires_posix(self) -> None: + from unittest import mock + + from pyspark.errors import PySparkRuntimeError + + with mock.patch.object(os, "name", "nt"): + with self.assertRaises(PySparkRuntimeError) as ctx: + local_server_pool.acquire_pooled_local_connect_server("local[2]", {}) + self.assertIn("POSIX", str(ctx.exception)) + + def test_attendant_warmup_stops_when_member_is_claimed(self) -> None: + from unittest import mock + + uid = "warmup" + os.makedirs(self._directory.member_dir(uid)) + attendant = MemberAttendant(self._directory, uid, "local[2]", "fp") + member = PoolMember(self._server_data(15002, os.getpid(), token="secret")) + child = mock.Mock() + child.poll.return_value = None + + with mock.patch.object(subprocess, "Popen", return_value=child) as popen: + with mock.patch.object(os.path, "exists", return_value=False): + attendant._warm(member) + + child.kill.assert_called_once_with() + child.wait.assert_called_once_with() + command = popen.call_args.args[0] + self.assertEqual(command[-3:], ["--warm", "--url", "sc://localhost:15002"]) + self.assertEqual( + popen.call_args.kwargs["env"]["SPARK_CONNECT_AUTHENTICATE_TOKEN"], "secret" + ) + + def test_run_warmup_executes_fixed_queries_and_stops_session(self) -> None: + from unittest import mock + + session = mock.MagicMock() + with mock.patch( + "pyspark.sql.connect.session.SparkSession", return_value=session + ) as connect_session: + local_server_pool._run_warmup("sc://localhost:15002") + + connect_session.assert_called_once_with(connection="sc://localhost:15002") + self.assertEqual(session.sql.call_count, 3) + session.range.assert_called_once_with(100000) + session.stop.assert_called_once_with() + + +@unittest.skipIf( + not should_test_connect or is_remote_only(), + connect_requirement_message or "Requires JVM access to start local Connect servers", +) +@unittest.skipUnless(os.name == "posix", "the pool relies on the POSIX sbin scripts") +class LocalConnectServerPoolE2ETests(unittest.TestCase): + """End-to-end tests that boot real pooled servers (slow).""" + + CLIENT = textwrap.dedent( + """ + import json + + from pyspark.sql import SparkSession + + spark = ( + SparkSession.builder.remote("local[2]") + .config("spark.local.connect.pool", "true") + .getOrCreate() + ) + from pyspark.sql.connect import local_server_pool + + try: + # Pool cleanup belongs only to the builder-created session that claimed the + # server. Stopping another session must leave the claimed server available. + spark.newSession().stop() + count = spark.range(1).count() + member = local_server_pool._claimed_member + print(json.dumps({"count": count, "server_pid": member.pid})) + finally: + spark.stop() + """ + ) + + def setUp(self) -> None: + self._tmpdir = tempfile.mkdtemp() + self._pool_dir = os.path.join(self._tmpdir, "pool") + self._saved_env = {k: os.environ.get(k) for k in _SAVED_ENV_KEYS} + os.environ["SPARK_LOCAL_CONNECT_POOL_DIR"] = self._pool_dir + + def tearDown(self) -> None: + try: + local_server_pool.purge_local_connect_pool() + deadline = time.time() + 60 + while time.time() < deadline: + if set(os.listdir(self._pool_dir)) <= {".lock"}: + break + # Supervising attendants notice the purged directory and exit on their own; + # purge again in case one republished state in between. + local_server_pool.purge_local_connect_pool() + time.sleep(1) + finally: + for k, v in self._saved_env.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + shutil.rmtree(self._tmpdir, ignore_errors=True) + + def _run_clients(self, n: int, pool_size: str) -> list: + env = dict(os.environ) + env["SPARK_LOCAL_CONNECT_POOL_DIR"] = self._pool_dir + env["SPARK_LOCAL_CONNECT_POOL"] = "1" + env["SPARK_LOCAL_CONNECT_POOL_SIZE"] = pool_size + env.pop("SPARK_CONNECT_AUTHENTICATE_TOKEN", None) + procs = [ + subprocess.Popen( + [sys.executable, "-c", self.CLIENT], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + for _ in range(n) + ] + outputs = [] + try: + for proc in procs: + stdout, stderr = proc.communicate(timeout=300) + 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() + return outputs + + def test_sequential_runs_use_fresh_servers_and_tear_them_down(self) -> None: + first = self._run_clients(1, pool_size="1")[0] + self.assertEqual(first["count"], 1) + # The used server is torn down asynchronously after its run. + self.assertTrue( + _wait_pid_gone(first["server_pid"]), + "the server was not torn down after its run ended", + ) + second = self._run_clients(1, pool_size="1")[0] + self.assertEqual(second["count"], 1) + self.assertNotEqual( + first["server_pid"], second["server_pid"], "a pooled server was reused across runs" + ) + + def test_concurrent_cold_clients_get_distinct_servers(self) -> None: + outputs = self._run_clients(2, pool_size="2") + self.assertEqual({o["count"] for o in outputs}, {1}) + pids = [o["server_pid"] for o in outputs] + self.assertEqual(len(set(pids)), 2, "two concurrent runs shared a pooled server") + + +if __name__ == "__main__": + from pyspark.testing import main + + main()