From bb5c0dd5af920e2ced836bfb77f766840b8b51ec Mon Sep 17 00:00:00 2001 From: ericm-db Date: Mon, 10 Aug 2026 18:43:34 +0000 Subject: [PATCH 1/7] [SPARK-58021][CONNECT] Add local server pool member claiming --- .../pyspark/sql/connect/local_server_pool.py | 150 ++++++++++- .../connect/test_connect_local_server_pool.py | 241 +++++++++++++++++- 2 files changed, 385 insertions(+), 6 deletions(-) diff --git a/python/pyspark/sql/connect/local_server_pool.py b/python/pyspark/sql/connect/local_server_pool.py index 72b13d5b52f48..fb5e1516b00f3 100644 --- a/python/pyspark/sql/connect/local_server_pool.py +++ b/python/pyspark/sql/connect/local_server_pool.py @@ -15,17 +15,136 @@ # limitations under the License. # -"""Filesystem-backed state storage for local Spark Connect server pools. +"""Filesystem-backed state model for local Spark Connect server pools. -This internal foundation owns the pool directory layout, locking, and JSON state-file access. +This internal foundation owns member identity, directory locking, state-file access, and +claiming. Process lifecycle and server acquisition are layered on top in follow-up changes. """ import contextlib +import hashlib import json +import math import os import shutil +import socket +import sys from typing import Any, Dict, List, Optional, Tuple +from pyspark.errors import PySparkValueError + + +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 + + +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 PySparkValueError(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 PySparkValueError("port is out of range") + if not math.isfinite(normalized["created"]): + raise PySparkValueError("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. @@ -229,3 +348,30 @@ def remove(self, path: str) -> None: def remove_member_dir(self, uid: str) -> None: self._assert_locked() shutil.rmtree(self.member_dir(uid), ignore_errors=True) + + +class ServerPool: + """Claims members from one pool directory; lifecycle operations are added later.""" + + 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 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 index 7e8d2326fbf8f..fe0756809c08e 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py @@ -15,20 +15,65 @@ # limitations under the License. # +import contextlib import os import shutil import subprocess import sys import tempfile +import time import unittest from pyspark.testing.connectutils import should_test_connect, connect_requirement_message if should_test_connect: - from pyspark.sql.connect.local_server_pool import PoolDirectory + 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__ -_SAVED_ENV_KEYS = ("SPARK_LOCAL_CONNECT_POOL_DIR",) +@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, + ) + + +_SAVED_ENV_KEYS = ( + "SPARK_LOCAL_CONNECT_POOL_DIR", + "PYSPARK_DRIVER_PYTHON", + "PYSPARK_PYTHON", +) # These tests start no server and exercise only stdlib filesystem code, so they do not need JVM @@ -40,7 +85,7 @@ ) @unittest.skipUnless(os.name == "posix", "the pool relies on POSIX file locks") class LocalConnectServerPoolUnitTests(unittest.TestCase): - """Tests for the pool filesystem storage; no real servers are started.""" + """Tests for the pool filesystem model; no real servers are started.""" def setUp(self) -> None: self._tmpdir = tempfile.mkdtemp() @@ -49,8 +94,16 @@ def setUp(self) -> None: 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 = [] def tearDown(self) -> 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) @@ -58,6 +111,33 @@ def tearDown(self) -> None: 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 _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") @@ -68,7 +148,7 @@ def test_pool_directory_location(self) -> None: 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") + state_path = self._directory.server_path("abcdef") with self._directory as directory: self.assertEqual(os.stat(directory.path).st_mode & 0o777, 0o700) lock_path = os.path.join(directory.path, ".lock") @@ -222,6 +302,159 @@ def test_not_reentrant(self) -> None: with self._directory: pass + @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 (("aaaa", time.time()), ("bbbb", 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-bbbb") + + 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 = { + "a0": {"fingerprint": "fp"}, + "a1": bad_pid, + "a2": bad_port, + "a3": bad_created, + "a4": non_finite_created, + "a5": out_of_range_port, + "a6": bad_host, + "a7": self._server_data( + port, sleeper.pid, spark_version="not-this-version" + ), + "a8": self._server_data(port, 2**31 - 1), + "a9": 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("b0"), + 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"}) + if __name__ == "__main__": from pyspark.testing import main From 893aa4b3e981f45367d0d39c1df80c3f9ec285e4 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Wed, 12 Aug 2026 17:33:21 +0000 Subject: [PATCH 2/7] [SPARK-58021][CONNECT] Apply Python formatting after rebase --- .../sql/tests/connect/test_connect_local_server_pool.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) 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 index fe0756809c08e..7c82ed3974802 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py @@ -352,9 +352,7 @@ def test_fingerprint_identity(self) -> None: # 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"}) - ) + 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"}) @@ -434,9 +432,7 @@ def test_claim_skips_malformed_and_incompatible_members(self) -> None: "a4": non_finite_created, "a5": out_of_range_port, "a6": bad_host, - "a7": self._server_data( - port, sleeper.pid, spark_version="not-this-version" - ), + "a7": self._server_data(port, sleeper.pid, spark_version="not-this-version"), "a8": self._server_data(port, 2**31 - 1), "a9": self._server_data(port, 2**100), } From 31d64ca385e3891a4a72a97be9e1d20b60b06533 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Wed, 12 Aug 2026 19:21:04 +0000 Subject: [PATCH 3/7] [SPARK-58021][CONNECT] Harden local server pool claiming --- .../pyspark/sql/connect/local_server_pool.py | 37 ++++-- .../connect/test_connect_local_server_pool.py | 117 ++++++++++++++++-- 2 files changed, 134 insertions(+), 20 deletions(-) diff --git a/python/pyspark/sql/connect/local_server_pool.py b/python/pyspark/sql/connect/local_server_pool.py index fb5e1516b00f3..34d866fcd6c7c 100644 --- a/python/pyspark/sql/connect/local_server_pool.py +++ b/python/pyspark/sql/connect/local_server_pool.py @@ -40,18 +40,23 @@ def pool_fingerprint(master: str, seed_conf: Dict[str, Any]) -> str: 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. + warehouse and Derby metastore locations resolve relative to it), the PySpark installation, + and the Python executable and path environment used for Python UDFs. """ server_python = os.environ.get( "PYSPARK_PYTHON", os.environ.get("PYSPARK_DRIVER_PYTHON", "python3") ) + resolved_server_python = shutil.which(server_python) or server_python + spark_home = os.environ.get("SPARK_HOME") identity = [ master, sorted((str(k), str(v)) for k, v in seed_conf.items()), os.getcwd(), sys.executable, - server_python, + resolved_server_python, + os.path.realpath(__file__), + os.path.realpath(spark_home) if spark_home else "", + os.environ.get("PYTHONPATH", ""), ] return hashlib.sha256(json.dumps(identity).encode("utf-8")).hexdigest()[:16] @@ -91,15 +96,22 @@ class PoolMember: 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 PySparkValueError(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: + if not isinstance(normalized[key], str) or not normalized[key]: + raise PySparkValueError(f"{key} must be a nonempty string") + for key in ("port", "pid"): + value = normalized[key] + if isinstance(value, bool) or not isinstance(value, int): + raise PySparkValueError(f"{key} must be an integer") + created = normalized["created"] + if isinstance(created, bool) or not isinstance(created, (int, float)): + raise PySparkValueError("created must be a number") + normalized["created"] = float(created) + if not 1 <= normalized["port"] <= 65535: raise PySparkValueError("port is out of range") - if not math.isfinite(normalized["created"]): - raise PySparkValueError("created must be finite") + if normalized["pid"] <= 0: + raise PySparkValueError("pid must be positive") + if not math.isfinite(normalized["created"]) or normalized["created"] < 0: + raise PySparkValueError("created must be finite and nonnegative") self.data = normalized # Set when this process claims the member; the path of its claimed--.json. self.claim_path: Optional[str] = None @@ -359,7 +371,8 @@ def __init__(self, directory: Optional[PoolDirectory] = None): 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.""" + rules use that pid to retire members whose client died without releasing them. The + caller must hold the directory lock so selection and rename form one transition.""" candidates = [] for uid, path in self._directory.paths_of_kind("server"): data = self._directory.read_json(path) 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 index 7c82ed3974802..70cbfbdde81c3 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py @@ -23,6 +23,7 @@ import tempfile import time import unittest +from unittest import mock from pyspark.testing.connectutils import should_test_connect, connect_requirement_message @@ -50,13 +51,14 @@ def _listening_socket(): listener.close() -def _closed_port() -> int: - """A port with nothing listening on it.""" +@contextlib.contextmanager +def _non_listening_socket(): + """Reserve a port without listening on it, so connection attempts are rejected.""" import socket with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.bind(("localhost", 0)) - return sock.getsockname()[1] + yield sock.getsockname()[1] def _spawn_sleeper() -> "subprocess.Popen": @@ -365,6 +367,59 @@ def test_fingerprint_identity(self) -> None: pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"}), ) + def test_fingerprint_resolves_server_python_from_path(self) -> None: + # Relative worker commands are resolved through PATH by the server. Include that + # resolution so equal command strings cannot identify different Python environments. + executable_name = "pool-test-python" + executable_dirs = [os.path.join(self._tmpdir, name) for name in ("env-a", "env-b")] + for directory in executable_dirs: + os.makedirs(directory) + executable = os.path.join(directory, executable_name) + with open(executable, "w") as executable_file: + executable_file.write("#!/bin/sh\n") + os.chmod(executable, 0o700) + os.environ["PYSPARK_PYTHON"] = executable_name + with mock.patch.dict(os.environ, {"PATH": executable_dirs[0]}): + first_environment = pool_fingerprint("local[*]", {}) + with mock.patch.dict(os.environ, {"PATH": executable_dirs[1]}): + self.assertNotEqual(first_environment, pool_fingerprint("local[*]", {})) + + def test_fingerprint_includes_python_and_spark_paths(self) -> None: + with mock.patch.dict(os.environ, {"PYTHONPATH": "/python/a", "SPARK_HOME": "/spark/a"}): + first_environment = pool_fingerprint("local[*]", {}) + with mock.patch.dict(os.environ, {"PYTHONPATH": "/python/b", "SPARK_HOME": "/spark/a"}): + self.assertNotEqual(first_environment, pool_fingerprint("local[*]", {})) + with mock.patch.dict(os.environ, {"PYTHONPATH": "/python/a", "SPARK_HOME": "/spark/b"}): + self.assertNotEqual(first_environment, pool_fingerprint("local[*]", {})) + + def test_pool_member_validation(self) -> None: + valid = self._server_data(12345, 123, created=1) + member = PoolMember.from_data(valid) + self.assertIsNotNone(member) + self.assertEqual(member.url, "sc://localhost:12345") + + invalid_records = { + "missing fields": {"fingerprint": "fp"}, + "empty token": self._server_data(12345, 123, token=""), + "non-string host": self._server_data(12345, 123, host=None), + "boolean port": self._server_data(True, 123), + "string port": self._server_data("12345", 123), + "fractional port": self._server_data(12345.5, 123), + "zero port": self._server_data(0, 123), + "out-of-range port": self._server_data(65536, 123), + "boolean pid": self._server_data(12345, True), + "string pid": self._server_data(12345, "123"), + "fractional pid": self._server_data(12345, 123.5), + "zero pid": self._server_data(12345, 0), + "string created": self._server_data(12345, 123, created="1"), + "boolean created": self._server_data(12345, 123, created=True), + "negative created": self._server_data(12345, 123, created=-1), + "non-finite created": self._server_data(12345, 123, created=float("nan")), + } + for name, data in invalid_records.items(): + with self.subTest(name=name): + self.assertIsNone(PoolMember.from_data(data)) + def test_claim_matches_fingerprint_and_renames(self) -> None: with _listening_socket() as port: sleeper = self._sleeper() @@ -402,12 +457,58 @@ def test_claim_prefers_the_oldest_member(self) -> None: # Prefer the oldest ready member for deterministic FIFO claiming. self.assertEqual(member.token, "t-bbbb") - def test_claim_skips_unreachable_member(self) -> None: - self._write_state( - self._directory.server_path("ccc"), self._server_data(_closed_port(), os.getpid()) + def test_concurrent_claimers_claim_one_member_once(self) -> None: + child = ( + "import sys\n" + "from pyspark.sql.connect.local_server_pool import PoolDirectory, ServerPool\n" + "directory = PoolDirectory(sys.argv[1])\n" + "with directory:\n" + " member = ServerPool(directory).claim('fp')\n" + "print(member.token if member is not None else 'NONE')\n" ) - with self._directory: - self.assertIsNone(self._pool.claim("fp")) + claimers = [] + results = [] + with _listening_socket() as port: + sleeper = self._sleeper() + uid = "cafe" + self._write_state( + self._directory.server_path(uid), + self._server_data(port, sleeper.pid, token="claimed-once"), + ) + try: + with self._directory: + for _ in range(2): + claimers.append( + subprocess.Popen( + [sys.executable, "-c", child, self._directory.path], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + ) + for proc in claimers: + stdout, stderr = proc.communicate(timeout=20) + self.assertEqual(proc.returncode, 0, stderr) + results.append(stdout.strip()) + finally: + for proc in claimers: + if proc.poll() is None: + proc.kill() + proc.communicate(timeout=10) + + self.assertEqual(sorted(results), ["NONE", "claimed-once"]) + states = self._states(uid) + self.assertEqual(set(states), {"claimed"}) + claiming_pid = PoolDirectory.claiming_pid(states["claimed"]) + self.assertIn(claiming_pid, [proc.pid for proc in claimers]) + + def test_claim_skips_unreachable_member(self) -> None: + with _non_listening_socket() as port: + self._write_state( + self._directory.server_path("ccc"), self._server_data(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: From 498a78f471a4dab02afb38bdf7acf352e80d8434 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Wed, 12 Aug 2026 19:29:58 +0000 Subject: [PATCH 4/7] [SPARK-58021][CONNECT] Store pool member state in fields --- .../pyspark/sql/connect/local_server_pool.py | 46 ++++++++----------- .../connect/test_connect_local_server_pool.py | 4 ++ 2 files changed, 22 insertions(+), 28 deletions(-) diff --git a/python/pyspark/sql/connect/local_server_pool.py b/python/pyspark/sql/connect/local_server_pool.py index 34d866fcd6c7c..f5116c1fd3abf 100644 --- a/python/pyspark/sql/connect/local_server_pool.py +++ b/python/pyspark/sql/connect/local_server_pool.py @@ -94,25 +94,31 @@ class PoolMember: """One published pool server, wrapping its ``server-.json`` record.""" def __init__(self, data: Dict[str, Any]): - normalized = dict(data) + record = dict(data) for key in ("host", "token", "spark_version", "fingerprint"): - if not isinstance(normalized[key], str) or not normalized[key]: + if not isinstance(record[key], str) or not record[key]: raise PySparkValueError(f"{key} must be a nonempty string") for key in ("port", "pid"): - value = normalized[key] + value = record[key] if isinstance(value, bool) or not isinstance(value, int): raise PySparkValueError(f"{key} must be an integer") - created = normalized["created"] + created = record["created"] if isinstance(created, bool) or not isinstance(created, (int, float)): raise PySparkValueError("created must be a number") - normalized["created"] = float(created) - if not 1 <= normalized["port"] <= 65535: + created = float(created) + if not 1 <= record["port"] <= 65535: raise PySparkValueError("port is out of range") - if normalized["pid"] <= 0: + if record["pid"] <= 0: raise PySparkValueError("pid must be positive") - if not math.isfinite(normalized["created"]) or normalized["created"] < 0: + if not math.isfinite(created) or created < 0: raise PySparkValueError("created must be finite and nonnegative") - self.data = normalized + self.host: str = record["host"] + self.port: int = record["port"] + self.token: str = record["token"] + self.pid: int = record["pid"] + self.spark_version: str = record["spark_version"] + self.fingerprint: str = record["fingerprint"] + self.created: float = created # Set when this process claims the member; the path of its claimed--.json. self.claim_path: Optional[str] = None @@ -124,36 +130,20 @@ def from_data(cls, data: Dict[str, Any]) -> Optional["PoolMember"]: 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']}" + return f"sc://{self.host}:{self.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): + if self.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 + return sock.connect_ex((self.host, self.port)) == 0 except (OSError, UnicodeError): return False 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 index 70cbfbdde81c3..8a40293ff84eb 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py @@ -396,6 +396,10 @@ def test_pool_member_validation(self) -> None: valid = self._server_data(12345, 123, created=1) member = PoolMember.from_data(valid) self.assertIsNotNone(member) + self.assertEqual(member.host, "localhost") + self.assertEqual(member.port, 12345) + self.assertEqual(member.pid, 123) + self.assertEqual(member.created, 1.0) self.assertEqual(member.url, "sc://localhost:12345") invalid_records = { From d6e5cc966fbfd60f86521ae386ee85977103f8d4 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Wed, 12 Aug 2026 19:39:10 +0000 Subject: [PATCH 5/7] [SPARK-58021][CONNECT] Clarify fingerprints and process tests --- .../pyspark/sql/connect/local_server_pool.py | 17 ++++-- .../connect/test_connect_local_server_pool.py | 61 ++++++++----------- 2 files changed, 38 insertions(+), 40 deletions(-) diff --git a/python/pyspark/sql/connect/local_server_pool.py b/python/pyspark/sql/connect/local_server_pool.py index f5116c1fd3abf..b6abfc4c2a20d 100644 --- a/python/pyspark/sql/connect/local_server_pool.py +++ b/python/pyspark/sql/connect/local_server_pool.py @@ -33,6 +33,8 @@ from pyspark.errors import PySparkValueError +_LINUX_ZOMBIE_STATE = "Z" + 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 @@ -58,7 +60,7 @@ def pool_fingerprint(master: str, seed_conf: Dict[str, Any]) -> str: os.path.realpath(spark_home) if spark_home else "", os.environ.get("PYTHONPATH", ""), ] - return hashlib.sha256(json.dumps(identity).encode("utf-8")).hexdigest()[:16] + return hashlib.sha256(json.dumps(identity).encode("utf-8")).hexdigest() def _pid_alive(pid: int) -> bool: @@ -78,15 +80,18 @@ def _pid_alive(pid: int) -> bool: 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() + with open(f"/proc/{pid}/status", encoding="utf-8") as status_file: + for line in status_file: + key, separator, value = line.partition(":") + if separator and key == "State": + state, _, _ = value.strip().partition(" ") + if state == _LINUX_ZOMBIE_STATE: + return False + break except FileNotFoundError: return False except OSError: pass - else: - if fields and fields[0] == "Z": - return False return True 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 index 8a40293ff84eb..b30e4222b6b7e 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py @@ -61,11 +61,11 @@ def _non_listening_socket(): yield sock.getsockname()[1] -def _spawn_sleeper() -> "subprocess.Popen": - """A long sleeper standing in for a pool server or attendant process.""" +def _spawn_live_process() -> "subprocess.Popen": + """A child blocked on its parent pipe, standing in for a live pool server.""" return subprocess.Popen( - [sys.executable, "-c", "import time; time.sleep(300)"], - stdin=subprocess.DEVNULL, + [sys.executable, "-c", "import sys; sys.stdin.buffer.read()"], + stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, ) @@ -113,8 +113,8 @@ def tearDown(self) -> None: os.environ[k] = v shutil.rmtree(self._tmpdir, ignore_errors=True) - def _sleeper(self) -> "subprocess.Popen": - proc = _spawn_sleeper() + def _live_process(self) -> "subprocess.Popen": + proc = _spawn_live_process() self._procs.append(proc) return proc @@ -305,22 +305,15 @@ def test_not_reentrant(self) -> None: pass @unittest.skipUnless( - sys.platform.startswith("linux") and os.path.isdir("/proc"), - "requires Linux process state", + sys.platform.startswith("linux") and os.path.isdir("/proc") and hasattr(os, "waitid"), + "requires Linux process state and waitid", ) 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") + # Wait for the child to exit but leave it waitable, which keeps it as a zombie + # until the finally block reaps it. + os.waitid(os.P_PID, proc.pid, os.WEXITED | os.WNOWAIT) self.assertFalse(local_server_pool._pid_alive(proc.pid)) finally: proc.wait(timeout=10) @@ -426,14 +419,14 @@ def test_pool_member_validation(self) -> None: def test_claim_matches_fingerprint_and_renames(self) -> None: with _listening_socket() as port: - sleeper = self._sleeper() + server_process = self._live_process() self._write_state( self._directory.server_path("aaa"), - self._server_data(port, sleeper.pid, fingerprint="other-fp"), + self._server_data(port, server_process.pid, fingerprint="other-fp"), ) self._write_state( self._directory.server_path("bbb"), - self._server_data(port, sleeper.pid, fingerprint="my-fp", token="t-bbb"), + self._server_data(port, server_process.pid, fingerprint="my-fp", token="t-bbb"), ) with self._directory: member = self._pool.claim("my-fp") @@ -450,11 +443,11 @@ def test_claim_matches_fingerprint_and_renames(self) -> None: def test_claim_prefers_the_oldest_member(self) -> None: with _listening_socket() as port: - sleeper = self._sleeper() + server_process = self._live_process() for uid, created in (("aaaa", time.time()), ("bbbb", time.time() - 100)): self._write_state( self._directory.server_path(uid), - self._server_data(port, sleeper.pid, token="t-" + uid, created=created), + self._server_data(port, server_process.pid, token="t-" + uid, created=created), ) with self._directory: member = self._pool.claim("fp") @@ -473,11 +466,11 @@ def test_concurrent_claimers_claim_one_member_once(self) -> None: claimers = [] results = [] with _listening_socket() as port: - sleeper = self._sleeper() + server_process = self._live_process() uid = "cafe" self._write_state( self._directory.server_path(uid), - self._server_data(port, sleeper.pid, token="claimed-once"), + self._server_data(port, server_process.pid, token="claimed-once"), ) try: with self._directory: @@ -516,18 +509,18 @@ def test_claim_skips_unreachable_member(self) -> None: 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) + server_process = self._live_process() + bad_pid = self._server_data(port, server_process.pid) bad_pid["pid"] = "not-a-pid" - bad_port = self._server_data(port, sleeper.pid) + bad_port = self._server_data(port, server_process.pid) bad_port["port"] = "not-a-port" - bad_created = self._server_data(port, sleeper.pid) + bad_created = self._server_data(port, server_process.pid) bad_created["created"] = "not-a-time" - non_finite_created = self._server_data(port, sleeper.pid) + non_finite_created = self._server_data(port, server_process.pid) non_finite_created["created"] = float("nan") - out_of_range_port = self._server_data(port, sleeper.pid) + out_of_range_port = self._server_data(port, server_process.pid) out_of_range_port["port"] = 65536 - bad_host = self._server_data(port, sleeper.pid) + bad_host = self._server_data(port, server_process.pid) bad_host["host"] = None records = { "a0": {"fingerprint": "fp"}, @@ -537,7 +530,7 @@ def test_claim_skips_malformed_and_incompatible_members(self) -> None: "a4": non_finite_created, "a5": out_of_range_port, "a6": bad_host, - "a7": self._server_data(port, sleeper.pid, spark_version="not-this-version"), + "a7": self._server_data(port, server_process.pid, spark_version="not-this-version"), "a8": self._server_data(port, 2**31 - 1), "a9": self._server_data(port, 2**100), } @@ -545,7 +538,7 @@ def test_claim_skips_malformed_and_incompatible_members(self) -> None: self._write_state(self._directory.server_path(uid), data) self._write_state( self._directory.server_path("b0"), - self._server_data(port, sleeper.pid, token="valid-token"), + self._server_data(port, server_process.pid, token="valid-token"), ) with self._directory: From 5e92368f37f01ba04e8b530ef087f5e779061525 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Wed, 12 Aug 2026 23:42:20 +0000 Subject: [PATCH 6/7] [SPARK-58021][CONNECT] Share liveness probes and complete pool fingerprint Address review feedback on the local Connect server pool foundation: - Consolidate _pid_alive and add a shared _port_open in local_server.py, so the reuse path (is_listening / is_reusable) and the pool path (is_usable) use the same hardened liveness and reachability probes. - Include both server-side Python interpreter resolutions in pool_fingerprint (SparkConnectPlanner.pythonExec prefers PYSPARK_PYTHON; PythonUtils. defaultPythonExec prefers PYSPARK_DRIVER_PYTHON), so a run differing only in PYSPARK_DRIVER_PYTHON no longer shares a member it would not have produced. - Add the JVM-shaping environment (SPARK_CONF_DIR, JAVA_HOME, etc.) to the fingerprint and soften the docstring to describe a curated, non-exhaustive set, since the launcher inherits the full environment. - Reject far-future created timestamps (beyond year 9999) so a corrupt value cannot look perpetually fresh to age-based reaping. - Document claim's lock-held blocking probe and its wall-clock ordering with sorted() tie-breaking. - Add tests: created inf/far-future/overflow, claim outside the lock, an already-claimed member being invisible, conf order-independence and str() keying, sys.executable, and JVM-env coverage. Co-authored-by: Isaac --- python/pyspark/sql/connect/local_server.py | 40 +++++- .../pyspark/sql/connect/local_server_pool.py | 124 ++++++++++-------- .../connect/test_connect_local_server_pool.py | 84 +++++++++++- 3 files changed, 184 insertions(+), 64 deletions(-) diff --git a/python/pyspark/sql/connect/local_server.py b/python/pyspark/sql/connect/local_server.py index 5ef479ee41c56..dc51415f4a13e 100644 --- a/python/pyspark/sql/connect/local_server.py +++ b/python/pyspark/sql/connect/local_server.py @@ -62,19 +62,53 @@ # A fixed SPARK_IDENT_STRING keeps the spark-daemon.sh pid and log file names stable # regardless of $USER. _SPARK_IDENT = "local-connect" +_LINUX_ZOMBIE_STATE = "Z" def _pid_alive(pid: int) -> bool: - """Whether ``pid`` exists (POSIX only). A process we cannot signal counts as alive.""" + """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 managed server. POSIX only, like the reuse and pool paths that call it. + """ + 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}/status", encoding="utf-8") as status_file: + for line in status_file: + key, separator, value = line.partition(":") + if separator and key == "State": + state, _, _ = value.strip().partition(" ") + if state == _LINUX_ZOMBIE_STATE: + return False + break + except FileNotFoundError: + return False + except OSError: + pass return True +def _port_open(host: str, port: int, timeout: float = 0.5) -> bool: + """Whether a TCP connection to ``host``:``port`` succeeds within ``timeout`` seconds. A + socket error or a host that fails to resolve counts as closed. + """ + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(timeout) + return sock.connect_ex((host, port)) == 0 + except (OSError, UnicodeError): + return False + + def _is_local_connect_server(pid: int) -> Optional[bool]: """Whether ``pid`` is still the managed Connect server recorded in discovery. @@ -223,9 +257,7 @@ def url(self) -> str: def is_listening(self) -> bool: if self.host is None or self.port is None: return False - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.settimeout(0.5) - return sock.connect_ex((self.host, self.port)) == 0 + return _port_open(self.host, self.port) def is_reusable(self) -> bool: from pyspark.version import __version__ diff --git a/python/pyspark/sql/connect/local_server_pool.py b/python/pyspark/sql/connect/local_server_pool.py index b6abfc4c2a20d..9c319cf98f633 100644 --- a/python/pyspark/sql/connect/local_server_pool.py +++ b/python/pyspark/sql/connect/local_server_pool.py @@ -27,72 +27,80 @@ import math import os import shutil -import socket import sys from typing import Any, Dict, List, Optional, Tuple from pyspark.errors import PySparkValueError -_LINUX_ZOMBIE_STATE = "Z" +# Environment variables that shape the JVM the launcher boots through +# sbin/start-connect-server.sh -> spark-daemon.sh -> load-spark-env.sh / spark-submit. +# SPARK_CONF_DIR selects the spark-env.sh and spark-defaults.conf that seed the server; the +# rest feed the classpath, heap, and JVM options. This is a curated set, not an exhaustive +# one: the launcher inherits the whole environment, so it names the inputs that most commonly +# differ between runs rather than every variable a server could read. +_JVM_ENV_VARS = ( + "SPARK_CONF_DIR", + "JAVA_HOME", + "SPARK_DIST_CLASSPATH", + "SPARK_DAEMON_MEMORY", + "SPARK_DRIVER_MEMORY", + "SPARK_SUBMIT_OPTS", + "SPARK_DAEMON_JAVA_OPTS", +) -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 PySpark installation, - and the Python executable and path environment used for Python UDFs. +def pool_fingerprint(master: str, seed_conf: Dict[str, Any]) -> str: + """The identity of a pool member: a curated set of inputs that shape 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. The set is curated + rather than complete because the launcher inherits the full environment (see + ``_JVM_ENV_VARS``) -- it covers the inputs that most commonly differ between runs. + + Besides the master and the seeded confs, this covers the working directory (unset warehouse + and Derby metastore locations resolve relative to it), the PySpark installation, the Python + interpreters the server would run UDFs and Python data sources with, and the environment + variables that shape the launched JVM. """ - server_python = os.environ.get( - "PYSPARK_PYTHON", os.environ.get("PYSPARK_DRIVER_PYTHON", "python3") + + def resolved(command: str) -> str: + # Relative commands resolve through PATH server-side; fold that in so equal command + # strings cannot stand for different interpreters on different PATHs. + return shutil.which(command) or command + + # Two server code paths resolve the Python interpreter with opposite precedence, so a run + # changing only one of these variables would still have booted a different server. Include + # both resolutions: SparkConnectPlanner.pythonExec (Connect Python UDFs) prefers + # PYSPARK_PYTHON, while PythonUtils.defaultPythonExec (Python data sources) prefers + # PYSPARK_DRIVER_PYTHON. Both fall back to python3 and treat an empty value as set, matching + # the Scala sys.env.getOrElse chains. + udf_python = resolved( + os.environ.get("PYSPARK_PYTHON", os.environ.get("PYSPARK_DRIVER_PYTHON", "python3")) + ) + data_source_python = resolved( + os.environ.get("PYSPARK_DRIVER_PYTHON", os.environ.get("PYSPARK_PYTHON", "python3")) ) - resolved_server_python = shutil.which(server_python) or server_python spark_home = os.environ.get("SPARK_HOME") identity = [ master, sorted((str(k), str(v)) for k, v in seed_conf.items()), os.getcwd(), sys.executable, - resolved_server_python, + udf_python, + data_source_python, os.path.realpath(__file__), os.path.realpath(spark_home) if spark_home else "", os.environ.get("PYTHONPATH", ""), + [os.environ.get(var, "") for var in _JVM_ENV_VARS], ] return hashlib.sha256(json.dumps(identity).encode("utf-8")).hexdigest() -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}/status", encoding="utf-8") as status_file: - for line in status_file: - key, separator, value = line.partition(":") - if separator and key == "State": - state, _, _ = value.strip().partition(" ") - if state == _LINUX_ZOMBIE_STATE: - return False - break - except FileNotFoundError: - return False - except OSError: - pass - return True +# The end of year 9999 UTC, as a Unix timestamp. ``created`` is a wall-clock ``time.time()`` +# reading, so no real clock reaches this for millennia; rejecting values beyond it keeps a +# corrupt far-future timestamp from looking perpetually fresh to age-based reaping in the +# layers above, which measure a member's age as ``time.time() - created``. +_MAX_CREATED = 253402300799 class PoolMember: @@ -115,8 +123,8 @@ def __init__(self, data: Dict[str, Any]): raise PySparkValueError("port is out of range") if record["pid"] <= 0: raise PySparkValueError("pid must be positive") - if not math.isfinite(created) or created < 0: - raise PySparkValueError("created must be finite and nonnegative") + if not math.isfinite(created) or not 0 <= created <= _MAX_CREATED: + raise PySparkValueError(f"created must be a finite timestamp in [0, {_MAX_CREATED}]") self.host: str = record["host"] self.port: int = record["port"] self.token: str = record["token"] @@ -140,17 +148,15 @@ def url(self) -> str: return f"sc://{self.host}:{self.port}" def is_usable(self) -> bool: - """Whether this member has a matching Spark version, live process, and open port.""" + """Whether this member has a matching Spark version, live process, and open port. Uses + the same liveness and reachability probes as the reuse path (see ``local_server``), so + the pool and reuse discovery agree on when a recorded server is still good.""" from pyspark.version import __version__ + from pyspark.sql.connect.local_server import _pid_alive, _port_open if self.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.host, self.port)) == 0 - except (OSError, UnicodeError): - return False + return _port_open(self.host, self.port) class PoolDirectory: @@ -367,7 +373,19 @@ 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. The - caller must hold the directory lock so selection and rename form one transition.""" + caller must hold the directory lock so selection and rename form one transition. + + Ordering is by ``created``, a wall-clock ``time.time()`` reading. It is comparable + across the independent processes that publish members, which ``time.monotonic()`` is + not, at the cost that a backward clock step (NTP, suspend/resume) can perturb the order. + Ties break by the stable ``sorted()`` over the sorted directory listing, so the order is + well defined but only approximately FIFO, not guaranteed. + + ``is_usable`` runs under the held lock and does blocking network I/O -- up to a 0.5s + connect for each candidate that is live but not accepting connections. The candidate + count is bounded by ``spark.local.connect.pool.size``, which is user-tunable, so a large + pool widens the window the lock is held; the reaping rules keep stale members from + accumulating without bound.""" candidates = [] for uid, path in self._directory.paths_of_kind("server"): data = self._directory.read_json(path) 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 index b30e4222b6b7e..495e409cda4d5 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py @@ -28,7 +28,7 @@ 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 import _pid_alive from pyspark.sql.connect.local_server_pool import ( PoolDirectory, PoolMember, @@ -314,7 +314,7 @@ def test_pid_alive_treats_zombie_as_dead(self) -> None: # Wait for the child to exit but leave it waitable, which keeps it as a zombie # until the finally block reaps it. os.waitid(os.P_PID, proc.pid, os.WEXITED | os.WNOWAIT) - self.assertFalse(local_server_pool._pid_alive(proc.pid)) + self.assertFalse(_pid_alive(proc.pid)) finally: proc.wait(timeout=10) @@ -343,8 +343,9 @@ def test_fingerprint_identity(self) -> None: 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. + # Match SparkConnectPlanner.pythonExec's PYSPARK_PYTHON -> PYSPARK_DRIVER_PYTHON -> + # python3 precedence for Connect UDFs, so clients never claim a server that resolved 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"})) @@ -352,10 +353,13 @@ def test_fingerprint_identity(self) -> None: self.assertNotEqual( base, pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"}) ) + # PYSPARK_DRIVER_PYTHON also feeds PythonUtils.defaultPythonExec (Python data sources), + # which prefers it over PYSPARK_PYTHON. So even with PYSPARK_PYTHON fixed, changing + # PYSPARK_DRIVER_PYTHON changes the server a run would boot and must change the identity. 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( + self.assertNotEqual( worker_python, pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"}), ) @@ -385,6 +389,38 @@ def test_fingerprint_includes_python_and_spark_paths(self) -> None: with mock.patch.dict(os.environ, {"PYTHONPATH": "/python/a", "SPARK_HOME": "/spark/b"}): self.assertNotEqual(first_environment, pool_fingerprint("local[*]", {})) + def test_fingerprint_conf_order_independent_and_string_keyed(self) -> None: + # sorted() over seed_conf makes the identity independent of dict insertion order, so a + # run that builds the same confs in a different order still matches. + self.assertEqual( + pool_fingerprint("local[*]", {"a": "1", "b": "2"}), + pool_fingerprint("local[*]", {"b": "2", "a": "1"}), + ) + # Confs serialize to a properties file, so values are compared as strings: 1 and "1" + # are the same seed and share an identity. + self.assertEqual( + pool_fingerprint("local[*]", {"k": 1}), pool_fingerprint("local[*]", {"k": "1"}) + ) + + def test_fingerprint_includes_python_executable(self) -> None: + # sys.executable is the client interpreter the server inherits; a packaging change that + # moves it must not silently reuse a server booted under the old one. + base = pool_fingerprint("local[*]", {}) + with mock.patch.object(sys, "executable", "/other/python"): + self.assertNotEqual(base, pool_fingerprint("local[*]", {})) + + def test_fingerprint_includes_jvm_env(self) -> None: + # Environment that shapes the launched JVM must change the identity. SPARK_CONF_DIR is + # the common CI case: it selects the spark-defaults.conf / spark-env.sh the server reads. + with mock.patch.dict(os.environ, {"SPARK_CONF_DIR": "/conf/a"}): + with_conf_a = pool_fingerprint("local[*]", {}) + with mock.patch.dict(os.environ, {"SPARK_CONF_DIR": "/conf/b"}): + self.assertNotEqual(with_conf_a, pool_fingerprint("local[*]", {})) + with mock.patch.dict(os.environ, {"JAVA_HOME": "/jdk/a"}): + with_java_a = pool_fingerprint("local[*]", {}) + with mock.patch.dict(os.environ, {"JAVA_HOME": "/jdk/b"}): + self.assertNotEqual(with_java_a, pool_fingerprint("local[*]", {})) + def test_pool_member_validation(self) -> None: valid = self._server_data(12345, 123, created=1) member = PoolMember.from_data(valid) @@ -411,7 +447,14 @@ def test_pool_member_validation(self) -> None: "string created": self._server_data(12345, 123, created="1"), "boolean created": self._server_data(12345, 123, created=True), "negative created": self._server_data(12345, 123, created=-1), - "non-finite created": self._server_data(12345, 123, created=float("nan")), + "nan created": self._server_data(12345, 123, created=float("nan")), + "infinite created": self._server_data(12345, 123, created=float("inf")), + # A finite float, but far past any real clock: rejected so it cannot look + # perpetually fresh to age-based reaping. + "far-future created": self._server_data(12345, 123, created=2**100), + # Too large to convert to float at all -- the OverflowError guard in from_data keeps + # a corrupt state file (this round-trips through json) from crashing the caller. + "overflow created": self._server_data(12345, 123, created=10**400), } for name, data in invalid_records.items(): with self.subTest(name=name): @@ -451,9 +494,32 @@ def test_claim_prefers_the_oldest_member(self) -> None: ) with self._directory: member = self._pool.claim("fp") - # Prefer the oldest ready member for deterministic FIFO claiming. + # Prefer the oldest ready member. Ordering is by wall-clock created, so this is + # approximate FIFO rather than a guarantee (see ServerPool.claim). self.assertEqual(member.token, "t-bbbb") + def test_claim_requires_the_lock(self) -> None: + # claim reaches the directory only through the locked accessors, so it inherits their + # assertion; pin the caller obligation directly so a future reordering that touches the + # directory before the first locked accessor is still caught. + with self.assertRaisesRegex(AssertionError, "context manager"): + self._pool.claim("fp") + + def test_claim_ignores_already_claimed_member(self) -> None: + # The kind filter is the exclusion invariant: a member already renamed to claimed-* is + # no longer of kind "server", so claim never hands it out a second time. A live, usable + # record under a claimed-* name must still be invisible to a new claimer. + with _listening_socket() as port: + server_process = self._live_process() + uid = "feed" + self._write_state( + self._directory.claimed_path(os.getpid() + 1, uid), + self._server_data(port, server_process.pid), + ) + with self._directory: + self.assertIsNone(self._pool.claim("fp")) + self.assertEqual(set(self._states(uid)), {"claimed"}) + def test_concurrent_claimers_claim_one_member_once(self) -> None: child = ( "import sys\n" @@ -493,6 +559,10 @@ def test_concurrent_claimers_claim_one_member_once(self) -> None: proc.kill() proc.communicate(timeout=10) + # Pins "claimed at most once": one child claims the member and the other sees none. It + # does not guarantee the two ever contend on the lock -- if the first child finishes + # before the second reaches flock(), the second simply finds no "server" entry -- so + # read this as a regression test for the exclusion invariant, not for lock contention. self.assertEqual(sorted(results), ["NONE", "claimed-once"]) states = self._states(uid) self.assertEqual(set(states), {"claimed"}) From 36c787ab5d6207eef8f876b87ad6b6b1688b7080 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Thu, 20 Aug 2026 19:06:07 +0000 Subject: [PATCH 7/7] [SPARK-58021][CONNECT] Address review: share the port probe, guard pid liveness in the helper, tighten the env test Address the second-round review on the local Connect server pool claiming layer: - ServerLauncher._await_ready now uses the shared _port_open helper instead of an inlined socket probe, finishing the reachability-probe consolidation and picking up its OSError/UnicodeError guard so a transient socket error retries until the deadline rather than aborting the launch. - Move the POSIX guard into _pid_alive itself (return True off POSIX, where os.kill would terminate the target rather than probe it) so every caller is safe; is_reusable drops its own os.name check. - Document PATH as a deliberate omission from the fingerprint's _JVM_ENV_VARS. - test_fingerprint_includes_jvm_env now covers all seven _JVM_ENV_VARS: it asserts the tuple matches an independent expected list (catching an accidental removal or an unlisted addition) and loops over each variable to prove it changes the identity. --- python/pyspark/sql/connect/local_server.py | 17 ++++++---- .../pyspark/sql/connect/local_server_pool.py | 6 ++++ .../connect/test_connect_local_server_pool.py | 34 +++++++++++++------ 3 files changed, 41 insertions(+), 16 deletions(-) diff --git a/python/pyspark/sql/connect/local_server.py b/python/pyspark/sql/connect/local_server.py index dc51415f4a13e..f8c0129f37752 100644 --- a/python/pyspark/sql/connect/local_server.py +++ b/python/pyspark/sql/connect/local_server.py @@ -68,8 +68,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 managed server. POSIX only, like the reuse and pool paths that call it. + or serve a managed server. + + Off POSIX this returns ``True`` without probing: ``os.kill`` there terminates the target for + any signal other than ``CTRL_C_EVENT`` / ``CTRL_BREAK_EVENT``, so signal 0 is not a safe + liveness probe, and callers fall through to the port check instead. Guarding here rather than + at each call site keeps every caller (reuse and pool) safe. (The pool path needs ``fcntl`` and + so never runs off POSIX regardless.) """ + if os.name != "posix": + return True if pid <= 0: return False try: @@ -264,7 +272,7 @@ def is_reusable(self) -> bool: if self.spark_version != __version__ or self.pid is None: return False - if os.name == "posix" and not _pid_alive(self.pid): + if not _pid_alive(self.pid): return False return self.is_listening() @@ -520,10 +528,7 @@ def _await_ready(self, port: int, token: str) -> None: while time.time() < deadline: pid = self._discovery.daemon_pid() if pid is not None: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.settimeout(0.5) - listening = sock.connect_ex(("localhost", port)) == 0 - if listening: + if _port_open("localhost", port): self._discovery.save( { "host": "localhost", diff --git a/python/pyspark/sql/connect/local_server_pool.py b/python/pyspark/sql/connect/local_server_pool.py index 9c319cf98f633..c506ab676abef 100644 --- a/python/pyspark/sql/connect/local_server_pool.py +++ b/python/pyspark/sql/connect/local_server_pool.py @@ -39,6 +39,12 @@ # rest feed the classpath, heap, and JVM options. This is a curated set, not an exhaustive # one: the launcher inherits the whole environment, so it names the inputs that most commonly # differ between runs rather than every variable a server could read. +# +# PATH is a deliberate omission: bin/spark-class prefers ${JAVA_HOME}/bin/java and only falls +# back to the first java on PATH, so two runs with different JDKs first on PATH and no JAVA_HOME +# would share a member. The identity already tracks PATH indirectly through shutil.which for the +# interpreters, and PATH is too volatile to fingerprint whole; a run needing a specific JDK +# should set JAVA_HOME. _JVM_ENV_VARS = ( "SPARK_CONF_DIR", "JAVA_HOME", 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 index 495e409cda4d5..48ea5970aa8db 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py @@ -30,6 +30,7 @@ if should_test_connect: from pyspark.sql.connect.local_server import _pid_alive from pyspark.sql.connect.local_server_pool import ( + _JVM_ENV_VARS, PoolDirectory, PoolMember, ServerPool, @@ -410,16 +411,29 @@ def test_fingerprint_includes_python_executable(self) -> None: self.assertNotEqual(base, pool_fingerprint("local[*]", {})) def test_fingerprint_includes_jvm_env(self) -> None: - # Environment that shapes the launched JVM must change the identity. SPARK_CONF_DIR is - # the common CI case: it selects the spark-defaults.conf / spark-env.sh the server reads. - with mock.patch.dict(os.environ, {"SPARK_CONF_DIR": "/conf/a"}): - with_conf_a = pool_fingerprint("local[*]", {}) - with mock.patch.dict(os.environ, {"SPARK_CONF_DIR": "/conf/b"}): - self.assertNotEqual(with_conf_a, pool_fingerprint("local[*]", {})) - with mock.patch.dict(os.environ, {"JAVA_HOME": "/jdk/a"}): - with_java_a = pool_fingerprint("local[*]", {}) - with mock.patch.dict(os.environ, {"JAVA_HOME": "/jdk/b"}): - self.assertNotEqual(with_java_a, pool_fingerprint("local[*]", {})) + # Every JVM-shaping variable must change the identity. SPARK_CONF_DIR is the common CI + # case (it selects the spark-defaults.conf / spark-env.sh the server reads); the rest + # feed the classpath, heap, and JVM options. The expected list is spelled out here + # independently rather than derived from _JVM_ENV_VARS: the fingerprint reads that same + # tuple, so dropping a variable from it would silently leave the fingerprint AND a loop + # over it in agreement. The equality check catches such drift (a removal or an unlisted + # addition), and the loop proves each variable still changes the identity. + expected = ( + "SPARK_CONF_DIR", + "JAVA_HOME", + "SPARK_DIST_CLASSPATH", + "SPARK_DAEMON_MEMORY", + "SPARK_DRIVER_MEMORY", + "SPARK_SUBMIT_OPTS", + "SPARK_DAEMON_JAVA_OPTS", + ) + self.assertEqual(set(_JVM_ENV_VARS), set(expected)) + for var in expected: + with self.subTest(var=var): + with mock.patch.dict(os.environ, {var: "/value/a"}): + with_a = pool_fingerprint("local[*]", {}) + with mock.patch.dict(os.environ, {var: "/value/b"}): + self.assertNotEqual(with_a, pool_fingerprint("local[*]", {})) def test_pool_member_validation(self) -> None: valid = self._server_data(12345, 123, created=1)