diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index 2faeea33b6f6..fd4a452d9750 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -1199,6 +1199,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_pool.py b/python/pyspark/sql/connect/local_server_pool.py new file mode 100644 index 000000000000..8bb4c5e09213 --- /dev/null +++ b/python/pyspark/sql/connect/local_server_pool.py @@ -0,0 +1,527 @@ +# +# 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 math +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_IDLE_TIMEOUT = 1800 + + +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), the client Python + executable, and the Python executable the Connect server selects for Python UDFs. + """ + server_python = os.environ.get( + "PYSPARK_PYTHON", os.environ.get("PYSPARK_DRIVER_PYTHON", "python3") + ) + identity = [ + master, + sorted((str(k), str(v)) for k, v in seed_conf.items()), + os.getcwd(), + sys.executable, + server_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 OverflowError: + 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, OverflowError): + return False + + +class PoolMember: + """One published pool server, wrapping its ``server-.json`` record.""" + + def __init__(self, data: Dict[str, Any]): + normalized = dict(data) + for key in ("host", "token", "spark_version", "fingerprint"): + if not isinstance(normalized[key], str): + raise ValueError(f"{key} must be a string") + normalized["port"] = int(normalized["port"]) + normalized["pid"] = int(normalized["pid"]) + normalized["created"] = float(normalized["created"]) + if not 0 <= normalized["port"] <= 65535: + raise ValueError("port is out of range") + if not math.isfinite(normalized["created"]): + raise ValueError("created must be finite") + self.data = normalized + # Set when this process claims the member; the path of its claimed--.json. + self.claim_path: Optional[str] = None + + @classmethod + def from_data(cls, data: Dict[str, Any]) -> Optional["PoolMember"]: + """Parse a published member record, returning ``None`` when it is malformed.""" + try: + return cls(data) + except (KeyError, TypeError, ValueError, OverflowError): + return 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["created"]) + + @property + def fingerprint(self) -> str: + return self.data["fingerprint"] + + @property + def url(self) -> str: + return f"sc://{self.data['host']}:{self.data['port']}" + + def is_usable(self) -> bool: + """Whether this member has a matching Spark version, live process, and open port.""" + from pyspark.version import __version__ + + if self.data["spark_version"] != __version__ or not _pid_alive(self.pid): + return False + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(0.5) + return sock.connect_ex((self.data["host"], self.data["port"])) == 0 + except (OSError, UnicodeError): + return False + + +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... + + Pool operations are infrequent, so one exclusive lock for every state transition is simpler + than a finer-grained scheme. A context can be entered again after exiting, allowing callers + to release the lock between polling attempts so other processes can update the directory. + """ + + 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) + # Re-assert privacy for an existing override directory: state files contain auth tokens, + # and directory write access would allow replacing them or bypassing the shared lock. + os.chmod(self.path, 0o700) + lock_fd = os.open(os.path.join(self.path, ".lock"), os.O_RDWR | os.O_CREAT, 0o600) + try: + os.fchmod(lock_fd, 0o600) + fcntl.flock(lock_fd, fcntl.LOCK_EX) + except BaseException: + os.close(lock_fd) + raise + self._lock_fd = lock_fd + 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: + """Claims, reaps, and retires 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) + member = PoolMember.from_data(data) if data is not None else None + if member is not None and member.fingerprint == fingerprint: + candidates.append((member, 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.from_data(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) + member = PoolMember.from_data(data) if data is not None else None + server_pid = member.pid if member is not None 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_pool.py b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py new file mode 100644 index 000000000000..a76680f08708 --- /dev/null +++ b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py @@ -0,0 +1,538 @@ +# +# 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_IDLE_TIMEOUT", + "PYSPARK_DRIVER_PYTHON", + "PYSPARK_PYTHON", +) + + +@unittest.skipIf( + not should_test_connect or is_remote_only(), + connect_requirement_message or "Requires Spark Connect test dependencies", +) +@unittest.skipUnless(os.name == "posix", "the pool relies on POSIX file locks") +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_directory_lock_and_state_file_permissions(self) -> None: + os.makedirs(self._directory.path, mode=0o755) + os.chmod(self._directory.path, 0o755) + state_path = self._directory.server_path("secure") + with self._directory as directory: + self.assertEqual(os.stat(directory.path).st_mode & 0o777, 0o700) + lock_path = os.path.join(directory.path, ".lock") + self.assertEqual(os.stat(lock_path).st_mode & 0o777, 0o600) + directory.write_json(state_path, {"token": "secret"}) + self.assertEqual(os.stat(state_path).st_mode & 0o777, 0o600) + self.assertEqual(directory.read_json(state_path), {"token": "secret"}) + + # O_CREAT does not apply its mode to an existing file, so writes must re-assert it. + os.chmod(state_path, 0o644) + directory.write_json(state_path, {"token": "new-secret"}) + self.assertEqual(os.stat(state_path).st_mode & 0o777, 0o600) + + with open(state_path, "w") as state_file: + state_file.write("not json") + with self._directory as directory: + self.assertIsNone(directory.read_json(state_path)) + + def test_pool_directory_lock_blocks_another_process(self) -> None: + child = ( + "import errno\n" + "import fcntl\n" + "import os\n" + "import sys\n" + "fd = os.open(sys.argv[1], os.O_RDWR)\n" + "try:\n" + " fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)\n" + "except OSError as error:\n" + " if error.errno not in (errno.EACCES, errno.EAGAIN):\n" + " raise\n" + "else:\n" + " raise RuntimeError('acquired a held lock')\n" + "finally:\n" + " os.close(fd)\n" + ) + with self._directory: + result = subprocess.run( + [sys.executable, "-c", child, os.path.join(self._directory.path, ".lock")], + capture_output=True, + text=True, + timeout=10, + ) + self.assertEqual(result.returncode, 0, result.stderr) + + @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"}) + ) + # Match SparkConnectPlanner's PYSPARK_PYTHON -> PYSPARK_DRIVER_PYTHON -> python3 + # precedence so clients never claim a server using a different fallback interpreter. + os.environ.pop("PYSPARK_PYTHON") + os.environ["PYSPARK_DRIVER_PYTHON"] = "python3" + self.assertEqual( + base, pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"}) + ) + os.environ["PYSPARK_DRIVER_PYTHON"] = "/driver/python" + self.assertNotEqual( + base, pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"}) + ) + os.environ["PYSPARK_PYTHON"] = "/worker/python" + worker_python = pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"}) + os.environ["PYSPARK_DRIVER_PYTHON"] = "/other/driver/python" + self.assertEqual( + worker_python, + 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_claim_skips_malformed_and_incompatible_members(self) -> None: + with _listening_socket() as port: + sleeper = self._sleeper() + bad_pid = self._server_data(port, sleeper.pid) + bad_pid["pid"] = "not-a-pid" + bad_port = self._server_data(port, sleeper.pid) + bad_port["port"] = "not-a-port" + bad_created = self._server_data(port, sleeper.pid) + bad_created["created"] = "not-a-time" + non_finite_created = self._server_data(port, sleeper.pid) + non_finite_created["created"] = float("nan") + out_of_range_port = self._server_data(port, sleeper.pid) + out_of_range_port["port"] = 65536 + bad_host = self._server_data(port, sleeper.pid) + bad_host["host"] = None + records = { + "missing": {"fingerprint": "fp"}, + "bad-pid": bad_pid, + "bad-port": bad_port, + "bad-created": bad_created, + "non-finite-created": non_finite_created, + "out-of-range-port": out_of_range_port, + "bad-host": bad_host, + "wrong-version": self._server_data( + port, sleeper.pid, spark_version="not-this-version" + ), + "dead": self._server_data(port, 2**31 - 1), + "overflowing-pid": self._server_data(port, 2**100), + } + for uid, data in records.items(): + self._write_state(self._directory.server_path(uid), data) + self._write_state( + self._directory.server_path("valid"), + self._server_data(port, sleeper.pid, token="valid-token"), + ) + + with self._directory: + member = self._pool.claim("fp") + + self.assertIsInstance(member, PoolMember) + self.assertEqual(member.token, "valid-token") + for uid in records: + self.assertEqual(set(self._states(uid)), {"server"}) + + 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()