This is an automated email from the ASF dual-hosted git repository.
dtenedor pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new ee11a92a9f1a [SPARK-58021][CONNECT] Add local server pool member
claiming
ee11a92a9f1a is described below
commit ee11a92a9f1a1f685b301b79d8267fb80e57d1e0
Author: Eric Marnadi <[email protected]>
AuthorDate: Fri Aug 21 13:20:07 2026 -0700
[SPARK-58021][CONNECT] Add local server pool member claiming
### What changes were proposed in this pull request?
This is layer 3 of the seven-PR local Connect pool stack:
#57684 -> #57685 -> #57907 -> #57686 -> #57687 -> #57102 -> #57688
The two lower layers are now merged, so GitHub shows only this layer's
two-file diff.
This layer adds member compatibility and claiming on top of the filesystem
state model:
- compatibility fingerprints covering the master, startup conf, working
directory, Python
executables, PySpark and Spark paths, and `PYTHONPATH`, using the full
SHA-256 digest;
- strict validation of persisted member records followed by explicit typed
`PoolMember` fields;
- Spark-version, process-liveness (including Linux zombies), and
socket-reachability checks; and
- well-defined, approximately-FIFO claiming through an atomic state-file
rename, including safe
concurrent claims from separate processes. Ordering is by the wall-clock
`created` timestamp
(comparable across the independent publisher processes, unlike
`time.monotonic()`) with stable
`sorted()` tie-breaking, so it is well defined but only approximately
FIFO, not guaranteed: a
backward clock step (NTP, suspend/resume) can perturb it.
Pool sizing lives with its first consumer in acquisition layer #57687.
Process lifecycle,
acquisition, SparkSession integration, and JIT warmup remain in later PRs.
### Why are the changes needed?
The filesystem layer defines safe state storage, but a client also needs to
distinguish compatible
servers and claim exactly one member without racing other processes.
Isolating that contract keeps
record validation and the ready-to-claimed transition independently
reviewable before lifecycle
and launch orchestration are added.
### Does this PR introduce _any_ user-facing change?
No. Claiming is internal and is not wired into SparkSession in this layer.
### How was this patch tested?
Added fifteen focused tests at this layer, bringing the suite to 25 tests.
They cover deterministic
Linux zombie detection, fingerprint identity and Python-interpreter
precedence, Python and Spark
path compatibility, JVM-environment coverage, strict member-record
validation, fingerprint-aware
and approximately-FIFO claiming, a real two-process claim race, unreachable
members, and malformed,
dead, or version-incompatible records.
```bash
python/run-tests --testnames
pyspark.sql.tests.connect.test_connect_local_server_pool
```
All 25 focused tests passed locally. Ruff lint and format checks, targeted
mypy, Python compilation,
custom-error validation, `git diff --check`, and changed-file ASCII and
line-length checks also
passed.
[GitHub Actions run
32406852918](https://github.com/ericm-db/spark/actions/runs/32406852918) is
running against the latest commit (`36c787ab5d6`).
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8) and OpenAI Codex (GPT-5)
Closes #57907 from ericm-db/local-connect-pool-claiming.
Authored-by: Eric Marnadi <[email protected]>
Signed-off-by: Daniel Tenedorio <[email protected]>
---
python/pyspark/sql/connect/local_server.py | 55 ++-
python/pyspark/sql/connect/local_server_pool.py | 182 ++++++++-
.../connect/test_connect_local_server_pool.py | 419 ++++++++++++++++++++-
3 files changed, 641 insertions(+), 15 deletions(-)
diff --git a/python/pyspark/sql/connect/local_server.py
b/python/pyspark/sql/connect/local_server.py
index 5ef479ee41c5..f8c0129f3775 100644
--- a/python/pyspark/sql/connect/local_server.py
+++ b/python/pyspark/sql/connect/local_server.py
@@ -62,19 +62,61 @@ _SERVER_CLASS =
"org.apache.spark.sql.connect.service.SparkConnectServer"
# 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.
+
+ 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:
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,16 +265,14 @@ class LocalConnectServer:
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__
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()
@@ -488,10 +528,7 @@ class ServerLauncher:
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 72b13d5b52f4..c506ab676abe 100644
--- a/python/pyspark/sql/connect/local_server_pool.py
+++ b/python/pyspark/sql/connect/local_server_pool.py
@@ -15,17 +15,155 @@
# 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 sys
from typing import Any, Dict, List, Optional, Tuple
+from pyspark.errors import PySparkValueError
+
+
+# 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.
+#
+# 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",
+ "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: 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.
+ """
+
+ 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"))
+ )
+ 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,
+ 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()
+
+
+# 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:
+ """One published pool server, wrapping its ``server-<uid>.json`` record."""
+
+ def __init__(self, data: Dict[str, Any]):
+ record = dict(data)
+ for key in ("host", "token", "spark_version", "fingerprint"):
+ 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 = record[key]
+ if isinstance(value, bool) or not isinstance(value, int):
+ raise PySparkValueError(f"{key} must be an integer")
+ created = record["created"]
+ if isinstance(created, bool) or not isinstance(created, (int, float)):
+ raise PySparkValueError("created must be a number")
+ created = float(created)
+ if not 1 <= record["port"] <= 65535:
+ raise PySparkValueError("port is out of range")
+ if record["pid"] <= 0:
+ raise PySparkValueError("pid must be positive")
+ 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"]
+ 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-<pid>-<uid>.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 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. 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
+ return _port_open(self.host, self.port)
+
class PoolDirectory:
"""Path layout, file access, and the cross-process lock of one pool
directory.
@@ -229,3 +367,43 @@ class PoolDirectory:
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-<pid>-<uid>.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.
+
+ 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)
+ 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 7e8d2326fbf8..48ea5970aa8d 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,68 @@
# limitations under the License.
#
+import contextlib
import os
import shutil
import subprocess
import sys
import tempfile
+import time
import unittest
+from unittest import mock
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.local_server import _pid_alive
+ from pyspark.sql.connect.local_server_pool import (
+ _JVM_ENV_VARS,
+ PoolDirectory,
+ PoolMember,
+ ServerPool,
+ pool_fingerprint,
+ )
+ from pyspark.version import __version__
-_SAVED_ENV_KEYS = ("SPARK_LOCAL_CONNECT_POOL_DIR",)
[email protected]
+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()
+
+
[email protected]
+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))
+ yield sock.getsockname()[1]
+
+
+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 sys; sys.stdin.buffer.read()"],
+ stdin=subprocess.PIPE,
+ 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 +88,7 @@ _SAVED_ENV_KEYS = ("SPARK_LOCAL_CONNECT_POOL_DIR",)
)
@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 +97,16 @@ class LocalConnectServerPoolUnitTests(unittest.TestCase):
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 +114,33 @@ class LocalConnectServerPoolUnitTests(unittest.TestCase):
os.environ[k] = v
shutil.rmtree(self._tmpdir, ignore_errors=True)
+ def _live_process(self) -> "subprocess.Popen":
+ proc = _spawn_live_process()
+ 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 +151,7 @@ class LocalConnectServerPoolUnitTests(unittest.TestCase):
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 +305,334 @@ class LocalConnectServerPoolUnitTests(unittest.TestCase):
with self._directory:
pass
+ @unittest.skipUnless(
+ 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:
+ # 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(_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.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"}))
+ os.environ["PYSPARK_DRIVER_PYTHON"] = "/driver/python"
+ 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.assertNotEqual(
+ worker_python,
+ 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_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:
+ # 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)
+ 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 = {
+ "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),
+ "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):
+ self.assertIsNone(PoolMember.from_data(data))
+
+ def test_claim_matches_fingerprint_and_renames(self) -> None:
+ with _listening_socket() as port:
+ server_process = self._live_process()
+ self._write_state(
+ self._directory.server_path("aaa"),
+ self._server_data(port, server_process.pid,
fingerprint="other-fp"),
+ )
+ self._write_state(
+ self._directory.server_path("bbb"),
+ self._server_data(port, server_process.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:
+ 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, server_process.pid, token="t-" +
uid, created=created),
+ )
+ with self._directory:
+ member = self._pool.claim("fp")
+ # 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"
+ "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"
+ )
+ claimers = []
+ results = []
+ with _listening_socket() as port:
+ server_process = self._live_process()
+ uid = "cafe"
+ self._write_state(
+ self._directory.server_path(uid),
+ self._server_data(port, server_process.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)
+
+ # 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"})
+ 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:
+ 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, server_process.pid)
+ bad_port["port"] = "not-a-port"
+ bad_created = self._server_data(port, server_process.pid)
+ bad_created["created"] = "not-a-time"
+ non_finite_created = self._server_data(port, server_process.pid)
+ non_finite_created["created"] = float("nan")
+ out_of_range_port = self._server_data(port, server_process.pid)
+ out_of_range_port["port"] = 65536
+ bad_host = self._server_data(port, server_process.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, server_process.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, server_process.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
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
