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/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..3fc0e9ec11dee --- /dev/null +++ b/python/pyspark/sql/connect/local_server_pool.py @@ -0,0 +1,509 @@ +# +# 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. +# + +"""Filesystem-backed state and lifecycle model for local Spark Connect server pools. + +This internal foundation owns member identity, directory locking, state-file access, claiming, +reaping, and retirement. Server acquisition is layered on top in a follow-up change. +""" + +import contextlib +import hashlib +import json +import os +import shutil +import signal +import socket +import sys +import time +from typing import Any, Dict, List, Optional, Tuple + +# 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 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 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 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. A later acquisition layer populates it; +# keeping the idempotent release path here makes lifecycle ownership explicit. +_claimed_member: Optional[PoolMember] = None + + +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`.""" + return ServerPool().purge() 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..492dcdf0f5ef8 --- /dev/null +++ b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py @@ -0,0 +1,438 @@ +# +# 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 os +import shutil +import subprocess +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 import local_server_pool + from pyspark.sql.connect.local_server_pool import ( + 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 + + +_SAVED_ENV_KEYS = ( + "SPARK_LOCAL_CONNECT_POOL_DIR", + "SPARK_LOCAL_CONNECT_POOL_SIZE", + "SPARK_LOCAL_CONNECT_POOL_IDLE_TIMEOUT", + "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) + + @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") + # Prefer the oldest ready member for deterministic FIFO claiming. + 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)) + + +if __name__ == "__main__": + from pyspark.testing import main + + main()