This is an automated email from the ASF dual-hosted git repository.
dtenedor pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 11d0c54a6607 [SPARK-58021][CONNECT] Generalize local Connect server
launching
11d0c54a6607 is described below
commit 11d0c54a6607ce8fa2977eae92490f30695fd6af
Author: Eric Marnadi <[email protected]>
AuthorDate: Mon Aug 10 10:03:44 2026 -0700
[SPARK-58021][CONNECT] Generalize local Connect server launching
### What changes were proposed in this pull request?
This is layer 1 of the six-PR local Connect pool stack:
#57684 -> #57685 -> #57686 -> #57687 -> #57102 -> #57688
This patch extracts the reusable pieces of local Spark Connect server
startup from the persistent-server reuse path:
- expose the per-user runtime directory and startup seed configuration as
module helpers;
- add `LocalConnectServer.start()` as the common lifecycle-owned launch
entry point; and
- let the launcher use an ephemeral port and a precomputed startup
configuration when requested.
The existing reuse path now calls the same `LocalConnectServer.start()`
method. Focused tests cover seed configuration and delegation of the new launch
options.
### Why are the changes needed?
The persistent reuse mode currently combines reusable local-server
lifecycle handling with assumptions specific to its one fixed daemon. The
single-use server pool needs the same secure configuration seeding, process
launch, discovery, and readiness handling, but with independent runtime
directories and ephemeral ports. Sharing one launch path keeps those behaviors
consistent and isolates the larger feature from the already-working reuse
implementation.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Added focused unit coverage to
`pyspark.sql.tests.connect.test_connect_local_server.LocalConnectServerReuseTests`
and ran the full suite, including real daemon startup, reuse, session
isolation, and static configuration seeding:
```bash
python -m unittest -v pyspark.sql.tests.connect.test_connect_local_server
```
All 15 tests passed. Ruff check, Ruff format check, and `git diff --check`
passed.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Fable 5) and OpenAI Codex (GPT-5)
Closes #57684 from ericm-db/local-connect-pool-refactor.
Authored-by: Eric Marnadi <[email protected]>
Signed-off-by: Daniel Tenedorio <[email protected]>
(cherry picked from commit 343fbc39f52fd3ee8bdcb125e52c683a0a90ab08)
Signed-off-by: Daniel Tenedorio <[email protected]>
---
python/pyspark/sql/connect/local_server.py | 146 +++++++++++-----
.../sql/tests/connect/test_connect_local_server.py | 184 +++++++++++++++++++++
2 files changed, 285 insertions(+), 45 deletions(-)
diff --git a/python/pyspark/sql/connect/local_server.py
b/python/pyspark/sql/connect/local_server.py
index 4f06d884aff4..4aa88a400487 100644
--- a/python/pyspark/sql/connect/local_server.py
+++ b/python/pyspark/sql/connect/local_server.py
@@ -75,6 +75,22 @@ def _pid_alive(pid: int) -> bool:
return True
+def runtime_dir() -> str:
+ """Return the private per-user directory holding local-server state."""
+ path = os.path.join(tempfile.gettempdir(),
"spark-connect-{}".format(getpass.getuser()))
+ try:
+ # exist_ok also covers two first runs racing to create the directory;
chmod
+ # re-asserts 0700 and fails if another user owns the path.
+ os.makedirs(path, mode=0o700, exist_ok=True)
+ os.chmod(path, 0o700)
+ except OSError as e:
+ raise PySparkRuntimeError(
+ errorClass="LOCAL_CONNECT_RUNTIME_DIR_UNAVAILABLE",
+ messageParameters={"path": path},
+ ) from e
+ return path
+
+
class Discovery:
"""Reads and writes the discovery file recording the persistent local
server.
@@ -82,26 +98,11 @@ class Discovery:
``SPARK_LOCAL_CONNECT_DISCOVERY`` points; the daemon's pid file and logs
sit next to it.
"""
- @staticmethod
- def _runtime_dir() -> str:
- path = os.path.join(tempfile.gettempdir(),
"spark-connect-{}".format(getpass.getuser()))
- try:
- # exist_ok also covers two first runs racing to create the
directory; chmod
- # re-asserts 0700 and fails if another user owns the path.
- os.makedirs(path, mode=0o700, exist_ok=True)
- os.chmod(path, 0o700)
- except OSError as e:
- raise PySparkRuntimeError(
- errorClass="LOCAL_CONNECT_RUNTIME_DIR_UNAVAILABLE",
- messageParameters={"path": path},
- ) from e
- return path
-
def __init__(self, path: Optional[str] = None):
self.path = os.path.abspath(
path
or os.environ.get("SPARK_LOCAL_CONNECT_DISCOVERY")
- or os.path.join(self._runtime_dir(), "connect-local.json")
+ or os.path.join(runtime_dir(), "connect-local.json")
)
self._lock_file: Optional[TextIO] = None
@@ -219,12 +220,33 @@ class LocalConnectServer:
def reuse_or_start(self, master: str, opts: Dict[str, Any]) -> str:
if not self.is_reusable():
- ServerLauncher(master, opts, self._discovery).launch()
- self._reload()
+ self.start(master, opts)
assert self.token is not None
os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = self.token
return self.url
+ def start(
+ self,
+ master: str,
+ opts: Dict[str, Any],
+ *,
+ use_ephemeral_port: bool = False,
+ seed_conf: Optional[Dict[str, Any]] = None,
+ ) -> None:
+ """Start this server and reload its discovery record.
+
+ Callers starting isolated daemons can request an ephemeral port and
provide a
+ precomputed startup configuration while sharing the standard launch
path.
+ """
+ ServerLauncher(
+ master,
+ opts,
+ self._discovery,
+ use_ephemeral_port=use_ephemeral_port,
+ seed_conf=seed_conf,
+ ).launch()
+ self._reload()
+
def stop(self) -> bool:
stopped = False
if self.pid is not None:
@@ -237,17 +259,58 @@ class LocalConnectServer:
return stopped
+def _strip_launcher_conf(conf: Dict[str, Any]) -> Dict[str, Any]:
+ """Drop keys the launcher sets itself (master, binding port, auth token)
and the
+ ``spark.local.connect.*`` opt-in keys, returning a new dict. Idempotent,
so it is safe to
+ apply to a conf that is already sanitized.
+ """
+ stripped = dict(conf)
+ for k in list(stripped):
+ if k in (
+ "spark.remote",
+ "spark.api.mode",
+ "spark.master",
+ "spark.connect.authenticate.token",
+ "spark.connect.grpc.binding.port",
+ ) or k.startswith("spark.local.connect."):
+ stripped.pop(k)
+ return stripped
+
+
+def startup_seed_conf(opts: Dict[str, Any]) -> Dict[str, Any]:
+ """Compute startup confs using the same merge as the in-process server
path, then strip
+ the keys the launcher manages (see ``_strip_launcher_conf``).
+ """
+ conf: Dict[str, Any] = {}
+ for i in range(int(os.environ.get("PYSPARK_REMOTE_INIT_CONF_LEN", "0"))):
+ conf = json.loads(os.environ["PYSPARK_REMOTE_INIT_CONF_{}".format(i)])
+ conf.update(opts)
+ return _strip_launcher_conf(conf)
+
+
class ServerLauncher:
"""Starts a persistent local server via ``sbin/start-connect-server.sh``
and waits until
it accepts connections. Callers must hold a ``Discovery`` context.
+
+ ``use_ephemeral_port`` and ``seed_conf`` allow other managed local servers
to share this
+ launch path without duplicating its process and readiness handling.
"""
_READY_TIMEOUT = 120
- def __init__(self, master: str, opts: Dict[str, Any], discovery:
Discovery):
+ def __init__(
+ self,
+ master: str,
+ opts: Dict[str, Any],
+ discovery: Discovery,
+ use_ephemeral_port: bool = False,
+ seed_conf: Optional[Dict[str, Any]] = None,
+ ):
self._master = master
self._opts = opts
self._discovery = discovery
+ self._use_ephemeral_port = use_ephemeral_port
+ self._seed_override = seed_conf
self._log_dir = os.path.join(discovery.directory, "logs")
def launch(self) -> None:
@@ -269,12 +332,12 @@ class ServerLauncher:
)
def _pick_port(self) -> int:
- """Under SPARK_TESTING always use an OS-assigned free port so suites
can run in
- parallel; otherwise honor the configured/default port, falling back to
a free one if
- another process holds it. (A live stale server of ours also holds the
port, but that
- start fails later at spark-daemon.sh's pid-file check regardless of
port.) The sbin
- script cannot report an ephemeral port back, so the free port is
picked and released
- here, with a small race until the server binds it.
+ """Use an OS-assigned free port when requested or under SPARK_TESTING
so suites can
+ run in parallel. Otherwise honor the configured/default port, falling
back to a free
+ one if another process holds it. (A live stale server of ours also
holds the port, but
+ that start fails later at spark-daemon.sh's pid-file check regardless
of port.) The
+ sbin script cannot report an ephemeral port back, so the free port is
picked and
+ released here, with a small race until the server binds it.
"""
def free_port() -> int:
@@ -282,7 +345,7 @@ class ServerLauncher:
sock.bind(("localhost", 0))
return sock.getsockname()[1]
- if "SPARK_TESTING" in os.environ:
+ if self._use_ephemeral_port or "SPARK_TESTING" in os.environ:
return free_port()
from pyspark.sql.connect.client import DefaultChannelBuilder
@@ -297,26 +360,19 @@ class ServerLauncher:
return free_port()
def _seed_conf(self) -> Dict[str, Any]:
- """Startup confs for the new server, merged like the in-process
- ``_start_connect_server`` merges ``PYSPARK_REMOTE_INIT_CONF_*`` with
the builder
- opts, minus the keys the launcher sets itself and the
``spark.local.connect.*``
- opt-in keys. Only the run that starts the server can seed static
confs; later runs
- find the JVM already warm.
+ """Startup confs for the new server, minus the keys the launcher sets
itself and the
+ ``spark.local.connect.*`` opt-in keys. Only the run that starts the
server can seed
+ static confs; later runs find the JVM already warm.
+
+ With no ``seed_conf`` override, this merges
``PYSPARK_REMOTE_INIT_CONF_*`` with the
+ builder opts like the in-process ``_start_connect_server`` does. When
an override is
+ given, it is used verbatim instead of the merge. Either way the result
is run through
+ ``_strip_launcher_conf``, so the launcher-managed keys never reach
+ ``--properties-file`` even if a caller passes raw opts as the override.
"""
- conf: Dict[str, Any] = {}
- for i in range(int(os.environ.get("PYSPARK_REMOTE_INIT_CONF_LEN",
"0"))):
- conf =
json.loads(os.environ["PYSPARK_REMOTE_INIT_CONF_{}".format(i)])
- conf.update(self._opts)
- for k in list(conf):
- if k in (
- "spark.remote",
- "spark.api.mode",
- "spark.master",
- "spark.connect.authenticate.token",
- "spark.connect.grpc.binding.port",
- ) or k.startswith("spark.local.connect."):
- conf.pop(k)
- return conf
+ if self._seed_override is not None:
+ return _strip_launcher_conf(self._seed_override)
+ return startup_seed_conf(self._opts)
@contextlib.contextmanager
def _seed_properties_file(self) -> Iterator[Optional[str]]:
diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server.py
b/python/pyspark/sql/tests/connect/test_connect_local_server.py
index eba26972ae8c..ae09eb04646f 100644
--- a/python/pyspark/sql/tests/connect/test_connect_local_server.py
+++ b/python/pyspark/sql/tests/connect/test_connect_local_server.py
@@ -104,6 +104,26 @@ class LocalConnectServerReuseTests(unittest.TestCase):
with Discovery() as discovery:
return LocalConnectServer(discovery)
+ def _launcher_discovery(self):
+ # A stand-in Discovery for ServerLauncher unit tests: only its
directory is read
+ # (for the log dir and the seed properties file), so point it at the
temp dir.
+ from unittest import mock
+
+ discovery = mock.Mock()
+ discovery.directory = self._tmpdir
+ return discovery
+
+ @contextlib.contextmanager
+ def _without_spark_testing(self):
+ # _pick_port's ephemeral branch is a no-op when SPARK_TESTING is set
(as it is under
+ # the test runner), so drop it to exercise the production behavior.
+ saved = os.environ.pop("SPARK_TESTING", None)
+ try:
+ yield
+ finally:
+ if saved is not None:
+ os.environ["SPARK_TESTING"] = saved
+
def test_discovery_location(self) -> None:
self.assertEqual(Discovery().path, self._discovery_path)
# Without the override the file lives in a per-user 0700 dir under the
temp dir.
@@ -114,6 +134,170 @@ class LocalConnectServerReuseTests(unittest.TestCase):
self.assertIn("spark-connect-{}".format(getpass.getuser()),
default.directory)
self.assertEqual(os.stat(default.directory).st_mode & 0o777, 0o700)
+ def test_startup_seed_conf(self) -> None:
+ from unittest import mock
+
+ initial = {
+ "spark.sql.shuffle.partitions": "8",
+ "spark.master": "local[1]",
+ }
+ opts = {
+ "spark.sql.warehouse.dir": os.path.join(self._tmpdir, "warehouse"),
+ "spark.local.connect.reuse": "true",
+ "spark.connect.grpc.binding.port": "0",
+ }
+ env = {
+ "PYSPARK_REMOTE_INIT_CONF_LEN": "1",
+ "PYSPARK_REMOTE_INIT_CONF_0": json.dumps(initial),
+ }
+ with mock.patch.dict(os.environ, env):
+ self.assertEqual(
+ local_server.startup_seed_conf(opts),
+ {
+ "spark.sql.shuffle.partitions": "8",
+ "spark.sql.warehouse.dir": opts["spark.sql.warehouse.dir"],
+ },
+ )
+
+ def test_start_delegates_launch_options(self) -> None:
+ from unittest import mock
+
+ discovery = mock.Mock()
+ discovery.load.side_effect = [
+ None,
+ {
+ "host": "localhost",
+ "port": 15002,
+ "token": "t",
+ "pid": os.getpid(),
+ "spark_version": __version__,
+ },
+ ]
+ server = LocalConnectServer(discovery)
+ seed_conf = {"spark.sql.shuffle.partitions": "4"}
+ with mock.patch.object(local_server, "ServerLauncher") as launcher:
+ server.start(
+ "local[2]",
+ {"spark.local.connect.reuse": "true"},
+ use_ephemeral_port=True,
+ seed_conf=seed_conf,
+ )
+
+ launcher.assert_called_once_with(
+ "local[2]",
+ {"spark.local.connect.reuse": "true"},
+ discovery,
+ use_ephemeral_port=True,
+ seed_conf=seed_conf,
+ )
+ launcher.return_value.launch.assert_called_once_with()
+ self.assertEqual(server.port, 15002)
+
+ def test_pick_port_uses_ephemeral_port_when_requested(self) -> None:
+ # This is the production path for pool attendants, which run without
SPARK_TESTING.
+ # A non-integer configured port would raise int() in the
configured/default branch;
+ # the ephemeral branch never reads it, so returning a clean
OS-assigned port proves
+ # the free-port path was taken even with SPARK_TESTING unset.
+ launcher = local_server.ServerLauncher(
+ "local[2]",
+ {"spark.local.connect.server.port": "not-a-port"},
+ self._launcher_discovery(),
+ use_ephemeral_port=True,
+ )
+ with self._without_spark_testing():
+ port = launcher._pick_port()
+ self.assertGreater(port, 0)
+
+ def test_pick_port_honors_configured_port_without_testing(self) -> None:
+ # With neither the ephemeral flag nor SPARK_TESTING, a free configured
port is used
+ # as-is rather than replaced by an OS-assigned one.
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+ sock.bind(("localhost", 0))
+ free = sock.getsockname()[1]
+ launcher = local_server.ServerLauncher(
+ "local[2]",
+ {"spark.local.connect.server.port": str(free)},
+ self._launcher_discovery(),
+ use_ephemeral_port=False,
+ )
+ with self._without_spark_testing():
+ self.assertEqual(launcher._pick_port(), free)
+
+ def test_seed_conf_override_is_used_and_sanitized(self) -> None:
+ # The override is taken verbatim except for launcher-managed keys,
which are stripped
+ # so raw builder opts passed as a seed cannot land in
--properties-file.
+ launcher = local_server.ServerLauncher(
+ "local[2]",
+ {},
+ self._launcher_discovery(),
+ seed_conf={
+ "spark.sql.shuffle.partitions": "4",
+ "spark.master": "local[9]",
+ "spark.local.connect.reuse": "true",
+ },
+ )
+ self.assertEqual(launcher._seed_conf(),
{"spark.sql.shuffle.partitions": "4"})
+
+ def test_seed_conf_empty_override_does_not_fall_through_to_env(self) ->
None:
+ from unittest import mock
+
+ # Load-bearing for the pool attendant: an empty override means "seed
nothing", and
+ # must not silently pick up PYSPARK_REMOTE_INIT_CONF_* the way
opts=None would.
+ env = {
+ "PYSPARK_REMOTE_INIT_CONF_LEN": "1",
+ "PYSPARK_REMOTE_INIT_CONF_0":
json.dumps({"spark.sql.shuffle.partitions": "8"}),
+ }
+ launcher = local_server.ServerLauncher(
+ "local[2]", {}, self._launcher_discovery(), seed_conf={}
+ )
+ with mock.patch.dict(os.environ, env):
+ self.assertEqual(launcher._seed_conf(), {})
+
+ def test_seed_conf_none_override_merges_env_and_opts(self) -> None:
+ from unittest import mock
+
+ # No override: the env-plus-opts merge (minus launcher-managed keys)
is used.
+ env = {
+ "PYSPARK_REMOTE_INIT_CONF_LEN": "1",
+ "PYSPARK_REMOTE_INIT_CONF_0":
json.dumps({"spark.sql.shuffle.partitions": "8"}),
+ }
+ launcher = local_server.ServerLauncher(
+ "local[2]",
+ {"spark.sql.warehouse.dir": os.path.join(self._tmpdir, "wh")},
+ self._launcher_discovery(),
+ seed_conf=None,
+ )
+ with mock.patch.dict(os.environ, env):
+ self.assertEqual(
+ launcher._seed_conf(),
+ {
+ "spark.sql.shuffle.partitions": "8",
+ "spark.sql.warehouse.dir": os.path.join(self._tmpdir,
"wh"),
+ },
+ )
+
+ def test_seed_properties_file_reflects_seed_conf(self) -> None:
+ # An empty seed yields no properties file, so start-connect-server.sh
gets no
+ # --properties-file; a non-empty seed writes a 0600 file with the
seeded confs.
+ launcher = local_server.ServerLauncher(
+ "local[2]", {}, self._launcher_discovery(), seed_conf={}
+ )
+ with launcher._seed_properties_file() as path:
+ self.assertIsNone(path)
+
+ launcher = local_server.ServerLauncher(
+ "local[2]",
+ {},
+ self._launcher_discovery(),
+ seed_conf={"spark.sql.shuffle.partitions": "4"},
+ )
+ with launcher._seed_properties_file() as path:
+ self.assertIsNotNone(path)
+ self.assertEqual(os.stat(path).st_mode & 0o777, 0o600)
+ with open(path) as f:
+ contents = f.read()
+ self.assertIn("spark.sql.shuffle.partitions=4", contents)
+
def test_discovery_roundtrip(self) -> None:
with Discovery() as discovery:
saved = self._server(port=15002)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
