[SPARK-57787][CONNECT] Reuse a persistent local Spark Connect server for faster local startup - #56907
[SPARK-57787][CONNECT] Reuse a persistent local Spark Connect server for faster local startup#56907ericm-db wants to merge 24 commits into
Conversation
|
Is this part of https://lists.apache.org/thread/sg9o2gbb3nttz74f0s01v8f167zy8ltt ? |
|
Oh yeah this was to address a comment that Nicholas had left on the doc he had linked in that thread. |
|
@nchammas I wasn't aware of the thread, only of the comment you left on the doc. Since you created the thread first - would you like me to close this PR so you can open what you had in mind? I'm totally fine either way. |
Review feedbackTwo issues in the persistent local Connect server path that are worth addressing before merge. 1. Concurrent first-time startup can fail spuriously (
|
If the committers prefer your approach, continue with your PR. For the record, what I concluded in that thread, after feedback from Holden and Tian, is that it would be better to take a different approach. My proposal was to create a new |
|
Thanks @dbtsai, both are real — fixed in df7b09a. 1. Concurrent first-time startup race. You're right, and it was masked in tests because
Verified with two processes launched simultaneously against the fixed port: both return the same endpoint, one starts and one reconnects, neither raises. Also verified the upgrade case falls back to an ephemeral port. 2. Builder options dropped on first startup. Also fixed. The starting client's confs (minus the keys the daemon controls itself — master/port/token/plugins and the reuse opt-in keys) are now merged the same way As you noted, this is inherently first-client-only: a later run reconnecting to an already-warm JVM can't change static confs. I've documented that limitation. |
Review feedback (follow-up)Timeout path can orphan the child JVM (
|
|
Good catch, fixed in bd938e7. The timeout path now signals the whole process group and escalates to SIGKILL if a graceful stop doesn't take, so the child JVM is reaped rather than orphaned: if os.name == "posix":
os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
else:
proc.terminate()
try:
proc.wait(timeout=10)
except Exception:
if os.name == "posix":
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
else:
proc.kill()This works because the daemon is a session leader ( I kept the daemon's own SIGTERM handler where it is (after |
|
Added a regression test for this in 321128d ( Also tightened the test |
|
@nchammas, I had misread the intention of your comment, I didn't think your proposal and this PR were in conflict, I thought we could proceed with both. |
Yes, we could proceed with both. There's no need to close this PR on my account! I have a working prototype of the unified CLI on my local branch. I am just waiting for some committer to show interest in the idea -- especially since it's a new user interface -- before pushing it and opening a PR. |
Hi @nchammas , I had an offline discussion with @ericm-db . We think you've put a lot of effort into the discussion and you probably want to start this project. Also the interface needs to be designed (if I think the idea of a new CLI interface is interesting, but you probably need more than my support. Anyway, I believe @ericm-db is also very interested in this project and you could work together on this! |
Do you think I need to create an SPIP, or shall I just open a PR and share it on the dev list? I wasn't sure if an SPIP would be overkill so I stuck to the dev list.
Happy to collaborate. I think a good sketch of what a collaboration might look like, given @ericm-db's work on this PR, is to figure out how automatic local discovery of an existing Spark Connect server would work. The unified CLI would be a relatively standalone piece of work. But as we make it easier for users to manage Connect servers via the new CLI, Eric can lead the effort to make leveraging that more seamless. In other words:
What do you think? |
|
@nchammas sounds great to me! |
|
Reopening after offline discussion with @nchammas |
|
Just to be clear to others reading along, our offline discussion was simply a reiteration of the plan I proposed earlier in this thread. I am waiting for feedback from @gaogaotiantian or some other committer on how best to proceed with regards to the unified CLI. |
dtenedor
left a comment
dtenedor
left a comment
There was a problem hiding this comment.
Thanks for working on this! Mostly LGTM, the test coverage is solid. We can figure this out on an opt-in basis for now while we think about next steps on the mailing list.
| pass | ||
|
|
||
| @staticmethod | ||
| def _stop_local_connect_server() -> bool: |
There was a problem hiding this comment.
This signals only the daemon PID, not the process group.
Unlike the timeout path, stop uses os.kill(pid, SIGTERM) only. If called while the daemon is still inside getOrCreate() (before its signal handler is installed), the JVM could be orphaned — same class of bug that was fixed for timeout. Unlikely in normal use (stop is a dev helper after the server is up), but inconsistent with _terminate_local_connect_server. Consider reusing the group-kill helper.
There was a problem hiding this comment.
Right, that makes sense. done
| every run and never leaks between runs. State backed by the shared `SparkContext` (the persistent | ||
| catalog/warehouse, global temp views, and cached datasets) *is* shared across runs, so namespace | ||
| per-run databases or clear that state yourself if your runs must be fully isolated. | ||
| - The server listens on port `15002` by default and authenticates with a token written, together |
There was a problem hiding this comment.
On version mismatch, a new server starts on an ephemeral port and writes discovery. The old server on 15002 keeps running until idle timeout. Over repeated upgrades, multiple idle JVMs could accumulate. Docs could mention killing the old pid from a stale discovery file.
| not should_test_connect or is_remote_only(), | ||
| connect_requirement_message or "Requires JVM access to start a local Connect server", | ||
| ) | ||
| class LocalConnectServerReuseTests(unittest.TestCase): |
There was a problem hiding this comment.
Could we also add an automated concurrent-startup test?
And one integration test through SparkSession.builder.remote("local[*]").getOrCreate() with the env var set?
|
@zhengruifeng @viirya , can you take a look as well? Thanks. |
nchammas
left a comment
nchammas
left a comment
There was a problem hiding this comment.
I am not a Connect expert, but it feels a bit off to be creating the discovery file from Python vs. via SparkConnectServer (the class that gets run when you call start-connect-server.sh) or SparkConnectService.
Ideally, we want any language client to be able to leverage the Connect server discovery
mechanism. Doesn't that mean the write/manage path for this mechanism should live in Scala?
| catalog/warehouse, global temp views, and cached datasets) *is* shared across runs, so namespace | ||
| per-run databases or clear that state yourself if your runs must be fully isolated. | ||
| - The server listens on port `15002` by default and authenticates with a token written, together | ||
| with the host, port, pid and Spark version, to `~/.spark/connect-local.json` (mode `0600`). |
There was a problem hiding this comment.
As far as I can tell, this is the first use of ~/.spark/ in the project. There are some aspects to discuss here since there has been some pushback against storing app data under ~/, even though it is very common.
On Linux, there is the XDG Base Directory Specification. On macOS, there is ~/Library/ and the various directories underneath. Windows probably has its own setup, too.
Python has the platformdirs library for taking care of these platform-specific differences. Perhaps Scala/Java also has something similar.
We should consider whether we want to try to follow these platform-specific conventions, even though our lives would be easier in the short term to just put things under ~/.spark/. Another idea is to use a specific path under SPARK_HOME. I think it would be a small win if we could eliminate the need for a new config or env var that's specific to this file's location.
There was a problem hiding this comment.
One hard constraint: the file holds a per-user auth token, so it must live in a per-user private location - that rules out SPARK_HOME, which is commonly shared/read-only (a system install serving multiple users would either collide or leak the token).
Within per-user options I don't feel strongly. ~/.spark/ follows the dev-tool convention Spark users already touch (~/.ivy2, ~/.m2, ~/.aws), but I'd be fine honoring $XDG_STATE_HOME/spark/ when that variable is set (no new dependency needed) and falling back to ~/.spark/ otherwise - platformdirs would mean a new runtime dep, which PySpark tries to avoid. SPARK_LOCAL_CONNECT_DISCOVERY already exists as the explicit override. If the Scala-side follow-up from your other comment happens, that's probably the right moment to settle the canonical path project-wide.
@nchammas Agreed on the end state: language-neutral discovery is the right goal. The natural design is SparkConnectService writing the discovery file itself when it binds (behind something like I kept the write path in Python here for scope reasons. The opt-in replaces a spawn path that's already Python-owned, and a Python component has to exist regardless, something client-side needs to make the reuse decision, launch the detached server on first use, and run the idle-timeout reaper. Moving the write into the server doesn't remove that piece, it just thins it: Python (and every other client) becomes a pure reader/launcher. My proposal: land this as the Python opt-in, then file a follow-up to move the write/manage path into SparkConnectService. That change touches the server and effectively defines the discovery file as a cross-language contract (format, location, multi-server behavior), so it deserves Connect-reviewer attention and probably belongs alongside the CLI SPIP discussion. Happy to file that JIRA and take it. If you'd rather settle the contract first, I'm open to that too, but I'd rather not grow this PR into a server change. |
gaogaotiantian
left a comment
gaogaotiantian
left a comment
There was a problem hiding this comment.
I think this implementation has plenty of details that need to be sorted out. But I want to talk about structures before we dig into details.
This whole feature is designed to be used locally for debug purpose. How safe do we want to make it? Do we need to protect from other users on the same machine? We might but that's a question.
We are saving about 2-3s per process spawn - that's a very important guidance. That's what time range we should work for. Having the process server persistent for 1 hr seems a bit too much for me.
I don't like how everything is implemented in sql/connect/session.py as static methods. Static methods should rarely be used - it basically means the method does not have to be in the class. As a matter of fact, I do believe most of the code should not be part of the class.
If you take a look at the current implementation, most of it is very specific to how local server is implemented. How the json data is stored, how the server deals with signals.
I believe most of the new session.py code should live in local_server.py, as a separate class. All the server specific code should be there (including how we pass the arguments). It should provide a single interface like class LocalConnectServer which handles all the life cycles (including spawning detached process that actually runs the server).
There should be no implementation detail of the local server itself in session.py. Let's write better Python code :)
| Once the server is accepting connections it writes a discovery file (host, the actually bound port, | ||
| the auth token, its pid and the Spark version) that later client processes read to reconnect. | ||
|
|
||
| It is launched by file path rather than ``python -m`` so it does not require the Spark Connect |
There was a problem hiding this comment.
This is feels wrong to me. A script in sql/connect that is executed by file - it does not fit the pattern. Why do we need to avoid requiring client requirements? This is already designed as a local debugging tool. I don't think we need to distinguish between "server" and "client". This script is executed by the user script which needs the client to work right?
There was a problem hiding this comment.
Agreed. Made the change
| # ordinary imports. Drop it before importing anything else (`import sys` cannot be shadowed); | ||
| # pyspark stays importable via the remaining sys.path entries. | ||
| if sys.path: | ||
| del sys.path[0] |
There was a problem hiding this comment.
Things like this is why I'm against running it as a file.
There was a problem hiding this comment.
Gone with the switch to use python -m instead of executing via file.
| def _write_discovery(path: str, host: str, port: int, token: str, version: str) -> None: | ||
| """Atomically write the discovery file with ``0600`` perms (it holds the auth token).""" | ||
| parent = os.path.dirname(path) | ||
| if parent and not os.path.isdir(parent): |
There was a problem hiding this comment.
What's this line protect against?
There was a problem hiding this comment.
Removed this line - os.mkdirs(parent, exist_ok = True) protects this.
| disc = json.load(f) | ||
| except (OSError, ValueError): | ||
| return | ||
| if disc.get("pid") == os.getpid(): |
There was a problem hiding this comment.
Do we care about any racing issues?
There was a problem hiding this comment.
Yeah there was one, between read and unlink.
A freshly started server could have replaced the file and the exiting daemon would delete the newcomer's entry.
It's fixed now though, the launching client has a lock, which removal needs to fetch. If the removal can't fetch it (startup is currently being executed), removal is skipped. Added a test for this.
| parser.add_argument( | ||
| "--idle-timeout", | ||
| type=float, | ||
| default=3600.0, |
There was a problem hiding this comment.
For local development, this seems a bit too long.
There was a problem hiding this comment.
Lowered to 1800s. Can tune further if you have something specific in mind.
| if disc.get("spark_version") != __version__: | ||
| return False | ||
| try: | ||
| os.kill(int(disc["pid"]), 0) |
There was a problem hiding this comment.
I don't think this is a no-op on Windows.
There was a problem hiding this comment.
yeah you're right - we need to skip on non POSIX systems. Added a test.
|
|
||
| It is launched by file path rather than ``python -m`` so it does not require the Spark Connect | ||
| *client* dependencies (grpc, etc.): a server only needs a classic PySpark install plus the Connect | ||
| server jar, like ``sbin/start-connect-server.sh``. It imports only the classic ``pyspark.sql`` API. |
There was a problem hiding this comment.
In order to do this, I think more correct way is directly invoke sbin/start-connect-server.sh instead of going through Py4J. The current way of using Py4J (without this change) itself is quite hacky
On the start-up timeout path, terminating only the detached daemon process left its child JVM orphaned when the timeout fired before the daemon had installed its own signal handling (the slow- startup case that triggers the timeout). Signal the whole process group -- the daemon is a session leader and the JVM stays in its group -- and escalate to SIGKILL if a graceful stop does not take, so the JVM is reaped instead of leaking.
Adds a POSIX-only test that starts the detached daemon, waits until it has spawned its child JVM, then calls _terminate_local_connect_server and asserts the whole process group is gone -- covering the orphaned-JVM case the timeout-path fix addresses. Also waits for the server port to close in tearDown so a stopped server's JVM cannot linger into the next test.
…ing the local server The stop path signals the pid recorded in the discovery file via killpg. If that pid is stale (e.g. the daemon was SIGKILLed and left its discovery file behind) and has been recycled by an unrelated process, group-killing it could take down the caller's own process group. The daemon is always launched as a session leader, so its group id equals its pid; only signal the group when that holds, and fall back to a plain kill otherwise. Adds a unit test covering both directions (with the signal syscalls mocked so a regression cannot kill the test runner), and a ruff-format fix for a missing blank line after the local_server.py module docstring. Co-authored-by: Isaac
…ves the daemon test_terminate_reaps_daemon_and_jvm fails in CI with "ProcessLookupError not raised": _terminate_local_connect_server keyed SIGKILL escalation on the daemon timing out, so when the daemon exits promptly on SIGTERM while its JVM's graceful shutdown lingers (or wedges), nothing ever escalates and the group survives. Terminate now waits for the whole process group to disappear after the daemon settles and group-SIGKILLs any survivors past a 10s grace period. It signals the group id directly since the leader may already be reaped by then, and polls the daemon so a zombie leader does not keep the group looking alive. Adds a deterministic regression test (a session leader that exits on SIGTERM after spawning a SIGTERM-immune child standing in for the slow JVM), and hardens both terminate tests against macOS zombie semantics: killpg on a group holding only a zombie yields EPERM, and kill(pid, 0) succeeds for a zombie until it is reaped. Co-authored-by: Isaac
…te tests Both terminate tests fail in CI with "process group survived": the tests checked group death via killpg(pgid, 0), but a zombie keeps its group signalable, and in CI containers whose pid 1 never reaps, the daemon's JVM reparents on the daemon's death and lingers as a zombie indefinitely -- so the check never turns negative even though nothing is running. (macOS launchd reaps promptly, which is why the tests pass locally.) The tests now assert that no *running* (non-zombie) process remains in the group, using ps to inspect member states. Co-authored-by: Isaac
…tion Review feedback: prefer fenced code blocks over Liquid highlight tags in the new documentation section; they render the same and keep raw-Markdown editors from mistaking comment hashes for headings. Co-authored-by: Isaac
… run the daemon via -m Addresses review comments: - Launch the daemon with `python -m pyspark.sql.connect.local_server` instead of by file path, dropping the sys.path workaround; the daemon always runs from the client's environment, which has the Connect client dependencies. - Move all client-side reuse helpers off SparkSession (they were all staticmethods) into pyspark/sql/connect/local_server.py as module functions; connect/session.py is back to its upstream state. - Pass the auth token only through SPARK_CONNECT_AUTHENTICATE_TOKEN (same precedence as the in-process path); it no longer appears on the daemon argv (visible in ps) nor as a duplicate conf (the server falls back to the env var when the conf is unset). - Probe pid liveness only on POSIX: os.kill(pid, 0) terminates the target process on Windows. - Close the remove-vs-publish race on the discovery file by claiming the start-up lock non-blockingly in _remove_discovery_if_ours. - Turn the start-up lock into a contextmanager; validate discovery value types on read and index the dict directly afterwards; simplify _discovery_path and the makedirs guards; drop the SIGINT try/except (main only runs on the main thread); reap the SIGKILLed daemon with a short wait instead of a bare poll. - Lower the default idle timeout to 1800s, add --stop to the module CLI, and update the docs accordingly. Co-authored-by: Isaac
…rt-connect-server.sh Reworks the opt-in local Connect server reuse path per review feedback: instead of a detached Python daemon that boots a classic Py4J session with the Connect plugin and babysits it, the launcher now drives the standard `sbin/start-connect-server.sh` (spark-daemon.sh submit SparkConnectServer) and lets spark-daemon.sh own the JVM (pid file, logs, daemonization). - No Python server process any more: PySpark runs the script with the auth token in its env, waits for the port, writes the discovery file (pid taken from the spark-daemon pid file), and returns the endpoint. - Deletes the Py4J usage, signal handling, idle-session polling, process-group management and SIGKILL escalation, and the discovery-file removal race guard -- none of it has an equivalent in the new shape. - Idle self-termination goes away with the daemon; the server runs until stopped (`python -m pyspark.sql.connect.local_server --stop`, or `sbin/stop-connect-server.sh`). `spark.local.connect.server.idleTimeout` is removed. - The server port is picked client-side (configured port when free, else an OS-assigned one), since the standalone script cannot report a bound ephemeral port back. - Seed confs travel via a 0600 spark-submit --properties-file instead of a JSON file parsed by the daemon. Artifact isolation needs no special handling: SparkConnectServer enables it itself. - The reuse path is POSIX-only now (it drives the sbin shell scripts) and says so explicitly; server logs land next to the discovery file (~/.spark/logs by default) for debuggability. - Ships start/stop-connect-server.sh in the pip sbin package (their dependencies, bin/* and spark-daemon.sh, were already packaged). Net effect vs the previous approach: -313 lines, and the server lifecycle is exactly the one users get from the documented manual workflow. Co-authored-by: Isaac
…e manual server workflow
Per review discussion, spark.local.connect.reuse / SPARK_LOCAL_CONNECT_REUSE
stays an internal, undocumented conf. The docs now cover only the explicit
workflow (sbin/start-connect-server.sh + .remote("sc://...")) plus a note on
what is and is not isolated between runs against a shared local server.
Co-authored-by: Isaac
Co-authored-by: Isaac
… move its state to a per-user temp dir Address review feedback on code structure, state location, and test strategy: - Rebuild local_server.py around three components instead of a flat list of module functions: LocalConnectServer (a typed record of one server with the operations scoped to it: url, is_reusable, stop), Discovery (the discovery file: location, load/save/clear, and the cross-process start-up lock), and ServerLauncher (port choice, seed confs, running the sbin script, readiness wait). The module entry points and the session.py integration are unchanged. - Move the discovery file, daemon pid file, and logs from ~/.spark to a per-user 0700 directory under the system temp dir, so nothing accumulates in permanent directories and stale records do not survive a reboot. The directory's ownership is verified before trusting it, and the module docstring now states the threat model explicitly. - Hold the start-up lock in stop_local_connect_server so a stop racing a concurrent launch cannot unlink the newcomer's record mid-start. - Rewrite the tests against the components' public interfaces instead of enumerating private functions: 19 tests become 13 with the same behavioral coverage, folding the reusability cases into subtests, dropping low-value tests of internal helpers, and keeping all four end-to-end tests. Co-authored-by: Isaac
Cut the module docstring down to what the code cannot say itself, drop docstrings that restated method names, fold the remaining comments into the few places with a non-obvious constraint (Windows os.kill semantics, token and conf handling off argv, the world-writable temp dir check), and remove the section banners from the test file. Co-authored-by: Isaac
…ry per review Address review feedback: document SPARK_LOCAL_CONNECT_REUSE in the Connect overview docs alongside the manual sbin workflow; normalize the discovery path so the getcwd/dirname special cases disappear; drop the atomic-rename dance in save() since all access holds the exclusive lock (re-asserting 0600 via fchmod); explain runtime-dir claim failures in the raised OSError; turn the seed properties file into a NamedTemporaryFile-backed context manager; move daemon pid reading into Discovery; and fix the stale upgrade story in the docs and the _pick_port docstring. Co-authored-by: Isaac Claude-Session: https://claude.ai/code/session_0146yRfjDdT8pWGFoR3Ydtez
Co-authored-by: Isaac Claude-Session: https://claude.ai/code/session_0146yRfjDdT8pWGFoR3Ydtez
…n the runtime dir cannot be claimed Co-authored-by: Isaac Claude-Session: https://claude.ai/code/session_019zdMT53oEXpAGcqmvqT1Yi
…ed server is stopped Co-authored-by: Isaac Claude-Session: https://claude.ai/code/session_019zdMT53oEXpAGcqmvqT1Yi
…erimental Co-authored-by: Isaac Claude-Session: https://claude.ai/code/session_019zdMT53oEXpAGcqmvqT1Yi
f4d9735 to
5f1f97b
Compare
…r cannot be claimed The PySpark custom errors check rejects raising the builtin PermissionError under python/pyspark/sql. Raise PySparkRuntimeError with a new LOCAL_CONNECT_RUNTIME_DIR_UNAVAILABLE condition instead, matching the rest of local_server.py. Co-authored-by: Isaac Claude-Session: https://claude.ai/code/session_01CFe4ZLGQWimrzCDdzvGMVU
|
@gaogaotiantian @dtenedor CI is green and this PR is ready to merge |
…for faster local startup
### What changes were proposed in this pull request?
Adds an opt-in fast path for local Spark Connect development. Today
`SparkSession.builder.remote("local[*]").getOrCreate()` starts a fresh in-process Connect server
that lives only as long as that Python process, so every run re-pays the cold start (JVM warmup,
`SparkContext` + server boot).
When `SPARK_LOCAL_CONNECT_REUSE=1` is set (or `spark.local.connect.reuse=true` on the builder), a
`local`-mode remote session instead reconnects to a persistent local Connect server. The first run
starts one through the standard `sbin/start-connect-server.sh` script (daemonized by
`sbin/spark-daemon.sh`, the same daemon a user would start by hand) and records host, port, auth
token, pid and Spark version in a discovery file. Later runs reuse that server if its Spark
version matches, its pid is alive, and its port accepts connections; otherwise they start a fresh
one. If a live but non-reusable server is still running (e.g. after a Spark upgrade), start-up
fails with an error pointing at the `--stop` command below, so old servers cannot silently
accumulate. User code is unchanged: the first run pays the cold start once, later runs reconnect
in a fraction of a second.
The implementation is `pyspark/sql/connect/local_server.py`, with `Discovery` handling the
discovery file (location, load/save, and the cross-process lock -- `Discovery` is a context
manager and all reads and writes happen under the lock), `LocalConnectServer` representing one
recorded server (its URL, reusability probe, stop), and `ServerLauncher` running the sbin script
and waiting for readiness. The first caller's startup confs are forwarded to the new server via a
`--properties-file`, mirroring what the in-process path does with its `SparkConf`.
The feature is off by default, Python-only, and POSIX-only (it relies on the `sbin/` shell
scripts). All state -- discovery file, daemon pid file, logs -- lives in a per-user `0700`
directory under the system temp dir rather than the home directory, so nothing accumulates
across reboots; `SPARK_LOCAL_CONNECT_DISCOVERY` overrides the location. The auth token is stored
with `0600` and the server binds localhost, so other users on the machine can neither read the
token nor reach the server. Each run is its own Connect session, so session-local state (temp
views, runtime SQL confs, artifacts) stays per-run; only shared `SparkContext` state (catalog,
global temp views, cached data) carries across runs.
The server runs until stopped with `python -m pyspark.sql.connect.local_server --stop`. It is an
ordinary `spark-daemon.sh` daemon, but it runs with a per-user pid directory and ident string so
it cannot collide with a manually started server -- which also means a plain
`sbin/stop-connect-server.sh` does not find it. The `--stop` command signals the recorded pid and
clears the discovery file; killing the pid directly also works, and the next run notices the dead
server and starts a fresh one. The earlier idle-timeout self-reaping was dropped with the switch
to the standard sbin daemon; a bounded lifetime would belong server-side (e.g. in
`SparkConnectService`), which also fits the follow-up discussed below of moving the
discovery-file write path into the server so any language client (and the proposed `spark
connect` CLI) can use the same mechanism.
### Why are the changes needed?
Creating a local Spark session for a quick edit/run loop takes a few seconds, and that cost is
one-time-per-process -- it does not amortize across separate runs. Keeping a warm server alive and
reconnecting to it is the only way to make a repeated local dev/test loop fast. This makes that
behavior available behind a single opt-in, without changing user code or default behavior.
### Does this PR introduce _any_ user-facing change?
Only when the opt-in is enabled. With `SPARK_LOCAL_CONNECT_REUSE=1` (or
`spark.local.connect.reuse=true` on the builder),
`SparkSession.builder.remote("local[*]").getOrCreate()` starts a persistent local Connect server on
the first run and reconnects on later runs, instead of booting a fresh in-process server each time.
With the opt-in unset (the default), behavior is unchanged. A new documentation section describes
the manual persistent-server workflow.
### How was this patch tested?
New `python/pyspark/sql/tests/connect/test_connect_local_server.py`. Unit tests cover the
`Discovery` save/load round-trip, file permissions and malformed-input rejection, the
`is_reusable` decision (version mismatch, dead pid, listening, closed port, and the Windows
pid-probe guard), stop semantics, the `--stop` CLI, and the POSIX guard. Four end-to-end tests
start a real server via the sbin scripts: builder integration with the opt-in, three concurrent
first-time startups converging on one server, reuse across calls with session isolation between
two connections, and the first caller's startup confs reaching the server's `SparkConf`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8, Fable 5)
Closes #56907 from ericm-db/local-connect-reuse.
Authored-by: Eric Marnadi <eric.marnadi@databricks.com>
Signed-off-by: Tian Gao <gaogaotiantian@hotmail.com>
(cherry picked from commit 0366e48)
Signed-off-by: Tian Gao <gaogaotiantian@hotmail.com>
|
Thank you for your patience @ericm-db ! |
|
Thanks for all the help and review @gaogaotiantian! |
6 participants
What changes were proposed in this pull request?
Adds an opt-in fast path for local Spark Connect development. Today
SparkSession.builder.remote("local[*]").getOrCreate()starts a fresh in-process Connect serverthat lives only as long as that Python process, so every run re-pays the cold start (JVM warmup,
SparkContext+ server boot).When
SPARK_LOCAL_CONNECT_REUSE=1is set (orspark.local.connect.reuse=trueon the builder), alocal-mode remote session instead reconnects to a persistent local Connect server. The first runstarts one through the standard
sbin/start-connect-server.shscript (daemonized bysbin/spark-daemon.sh, the same daemon a user would start by hand) and records host, port, authtoken, pid and Spark version in a discovery file. Later runs reuse that server if its Spark
version matches, its pid is alive, and its port accepts connections; otherwise they start a fresh
one. If a live but non-reusable server is still running (e.g. after a Spark upgrade), start-up
fails with an error pointing at the
--stopcommand below, so old servers cannot silentlyaccumulate. User code is unchanged: the first run pays the cold start once, later runs reconnect
in a fraction of a second.
The implementation is
pyspark/sql/connect/local_server.py, withDiscoveryhandling thediscovery file (location, load/save, and the cross-process lock --
Discoveryis a contextmanager and all reads and writes happen under the lock),
LocalConnectServerrepresenting onerecorded server (its URL, reusability probe, stop), and
ServerLauncherrunning the sbin scriptand waiting for readiness. The first caller's startup confs are forwarded to the new server via a
--properties-file, mirroring what the in-process path does with itsSparkConf.The feature is off by default, Python-only, and POSIX-only (it relies on the
sbin/shellscripts). All state -- discovery file, daemon pid file, logs -- lives in a per-user
0700directory under the system temp dir rather than the home directory, so nothing accumulates
across reboots;
SPARK_LOCAL_CONNECT_DISCOVERYoverrides the location. The auth token is storedwith
0600and the server binds localhost, so other users on the machine can neither read thetoken nor reach the server. Each run is its own Connect session, so session-local state (temp
views, runtime SQL confs, artifacts) stays per-run; only shared
SparkContextstate (catalog,global temp views, cached data) carries across runs.
The server runs until stopped with
python -m pyspark.sql.connect.local_server --stop. It is anordinary
spark-daemon.shdaemon, but it runs with a per-user pid directory and ident string soit cannot collide with a manually started server -- which also means a plain
sbin/stop-connect-server.shdoes not find it. The--stopcommand signals the recorded pid andclears the discovery file; killing the pid directly also works, and the next run notices the dead
server and starts a fresh one. The earlier idle-timeout self-reaping was dropped with the switch
to the standard sbin daemon; a bounded lifetime would belong server-side (e.g. in
SparkConnectService), which also fits the follow-up discussed below of moving thediscovery-file write path into the server so any language client (and the proposed
spark connectCLI) can use the same mechanism.Why are the changes needed?
Creating a local Spark session for a quick edit/run loop takes a few seconds, and that cost is
one-time-per-process -- it does not amortize across separate runs. Keeping a warm server alive and
reconnecting to it is the only way to make a repeated local dev/test loop fast. This makes that
behavior available behind a single opt-in, without changing user code or default behavior.
Does this PR introduce any user-facing change?
Only when the opt-in is enabled. With
SPARK_LOCAL_CONNECT_REUSE=1(orspark.local.connect.reuse=trueon the builder),SparkSession.builder.remote("local[*]").getOrCreate()starts a persistent local Connect server onthe first run and reconnects on later runs, instead of booting a fresh in-process server each time.
With the opt-in unset (the default), behavior is unchanged. A new documentation section describes
the manual persistent-server workflow.
How was this patch tested?
New
python/pyspark/sql/tests/connect/test_connect_local_server.py. Unit tests cover theDiscoverysave/load round-trip, file permissions and malformed-input rejection, theis_reusabledecision (version mismatch, dead pid, listening, closed port, and the Windowspid-probe guard), stop semantics, the
--stopCLI, and the POSIX guard. Four end-to-end testsstart a real server via the sbin scripts: builder integration with the opt-in, three concurrent
first-time startups converging on one server, reuse across calls with session isolation between
two connections, and the first caller's startup confs reaching the server's
SparkConf.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8, Fable 5)