[SPARK-58021][CONNECT] Add local server pool member claiming - #57907
[SPARK-58021][CONNECT] Add local server pool member claiming#57907ericm-db wants to merge 6 commits into
Conversation
### What changes were proposed in this pull request? This is layer 2 of the seven-PR local Connect pool stack: #57684 -> #57685 -> #57907 -> #57686 -> #57687 -> #57102 -> #57688 The review unit introduced here is commit `23f806b64ad`. This layer adds the filesystem-backed storage foundation for pool members: - stable state-file paths keyed by member ID and per-member directories; - an overridable private pool directory under the per-user runtime directory; - a per-pool cross-process POSIX file lock; - private directory, lock-file, and JSON state-file permissions; and - locked helpers for listing, reading, writing, renaming, and removing member state. Member validation, compatibility fingerprints, and atomic claiming are isolated in #57907. Process lifecycle, acquisition, SparkSession integration, and JIT warmup remain in later PRs. ### Why are the changes needed? The pool needs a small, independently reviewable state model before adding compatibility checks, claiming, and process supervision. Keeping this layer limited to path layout, locking, and state file access makes its filesystem and concurrency contract reviewable on its own. ### Does this PR introduce _any_ user-facing change? No. The storage model is internal and is not wired into SparkSession in this layer. ### How was this patch tested? Added three focused tests covering directory selection, private permissions and malformed JSON, and cross-process lock contention. ```bash python/run-tests --testnames pyspark.sql.tests.connect.test_connect_local_server_pool ``` These cases passed on Python 3.11 as part of the combined suite before the stack was split. The rebuilt commit passed `git diff --check`, Python AST parsing, and changed-line ASCII and 100-column checks. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Fable 5) and OpenAI Codex (GPT-5) Closes #57685 from ericm-db/local-connect-pool-storage. Authored-by: Eric Marnadi <eric.marnadi@databricks.com> Signed-off-by: Daniel Tenedorio <daniel.tenedorio@databricks.com>
### What changes were proposed in this pull request? This is layer 2 of the seven-PR local Connect pool stack: #57684 -> #57685 -> #57907 -> #57686 -> #57687 -> #57102 -> #57688 The review unit introduced here is commit `23f806b64ad`. This layer adds the filesystem-backed storage foundation for pool members: - stable state-file paths keyed by member ID and per-member directories; - an overridable private pool directory under the per-user runtime directory; - a per-pool cross-process POSIX file lock; - private directory, lock-file, and JSON state-file permissions; and - locked helpers for listing, reading, writing, renaming, and removing member state. Member validation, compatibility fingerprints, and atomic claiming are isolated in #57907. Process lifecycle, acquisition, SparkSession integration, and JIT warmup remain in later PRs. ### Why are the changes needed? The pool needs a small, independently reviewable state model before adding compatibility checks, claiming, and process supervision. Keeping this layer limited to path layout, locking, and state file access makes its filesystem and concurrency contract reviewable on its own. ### Does this PR introduce _any_ user-facing change? No. The storage model is internal and is not wired into SparkSession in this layer. ### How was this patch tested? Added three focused tests covering directory selection, private permissions and malformed JSON, and cross-process lock contention. ```bash python/run-tests --testnames pyspark.sql.tests.connect.test_connect_local_server_pool ``` These cases passed on Python 3.11 as part of the combined suite before the stack was split. The rebuilt commit passed `git diff --check`, Python AST parsing, and changed-line ASCII and 100-column checks. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Fable 5) and OpenAI Codex (GPT-5) Closes #57685 from ericm-db/local-connect-pool-storage. Authored-by: Eric Marnadi <eric.marnadi@databricks.com> Signed-off-by: Daniel Tenedorio <daniel.tenedorio@databricks.com> (cherry picked from commit 001e89c) Signed-off-by: Daniel Tenedorio <daniel.tenedorio@databricks.com>
06648f8 to
31d64ca
Compare
dtenedor
left a comment
There was a problem hiding this comment.
Review notes on correctness and test coverage. The claiming protocol itself looks right: selection and the server- -> claimed-<pid>- rename happen inside one flock(LOCK_EX) critical section, so the rename is the mutual exclusion, and a losing claimer simply stops seeing the entry as kind server. The two-process test confirms that against real processes rather than mocked locking. The items below are what I'd want addressed.
Correctness
1. _pid_alive and the reachability probe are near-duplicates of code already in the same package, with divergent semantics. local_server.py -- which this module already imports runtime_dir from -- has its own version:
def _pid_alive(pid: int) -> bool:
"""Whether ``pid`` exists (POSIX only). A process we cannot signal counts as alive."""
try:
os.kill(pid, 0)
except ProcessLookupError:
return False
except OSError:
pass
return TrueIt also has is_listening() (identical AF_INET / 0.5s / connect_ex == 0 logic, but with no except (OSError, UnicodeError) guard) and is_reusable(), which is structurally the same version-then-pid-then-port sequence as the new PoolMember.is_usable(). The new copies are strictly better -- the pid <= 0 guard, OverflowError, Linux zombies, and the exception guard around connect_ex are all absent from the old ones. That means the hardening lands only on the pool path while the reuse path keeps the weaker behavior, and a reader has no way to tell which definition is authoritative. I'd import _pid_alive from local_server (or lift both helpers to one place) so there is a single answer to "is this server alive".
2. The fingerprint drops PYSPARK_DRIVER_PYTHON whenever PYSPARK_PYTHON is set, but the server resolves the interpreter with the opposite precedence in another code path. Mirroring SparkConnectPlanner.pythonExec (PYSPARK_PYTHON -> PYSPARK_DRIVER_PYTHON -> python3) is right for Connect Python UDFs, and the comment is accurate. But PythonUtils.defaultPythonExec reverses it (PYSPARK_DRIVER_PYTHON -> PYSPARK_PYTHON -> python3), and it is what DataSourceManager.shouldLoadPythonDataSources gates on and what PythonUtils.createPythonFunction uses -- including deriving pythonVer by actually executing it. Both run inside the server. So two runs differing only in PYSPARK_DRIVER_PYTHON get the same fingerprint yet would have booted servers whose Python-data-source interpreter and version differ. The test currently asserts that equality:
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(...))which pins the gap as intended behavior. Including both raw values (or both resolutions) in the identity list costs nothing and removes the question.
3. Environment that shapes the launched JVM is missing from the fingerprint. local_server.py::_run_script hands env = dict(os.environ) to $SPARK_HOME/sbin/start-connect-server.sh, which runs through spark-daemon.sh and load-spark-env.sh -- where SPARK_CONF_DIR defaults to $SPARK_HOME/conf and sources spark-env.sh, and spark-submit reads spark-defaults.conf from that same directory. None of SPARK_CONF_DIR, JAVA_HOME, SPARK_DIST_CLASSPATH, SPARK_DAEMON_MEMORY / SPARK_DRIVER_MEMORY, or SPARK_SUBMIT_OPTS / SPARK_DAEMON_JAVA_OPTS appears in the identity. Two runs differing only in SPARK_CONF_DIR -- a common CI pattern -- would share a member whose server-side defaults came from someone else's conf directory. Either extend the list, or soften the docstring: "everything that shapes the server a run would have booted for itself" promises completeness that the implementation doesn't deliver, and the honest version ("a curated set, because the server inherits the full environment") is more useful to the next reader.
4. claim does blocking network I/O while holding the exclusive pool lock. Each candidate that is live but not accepting connections costs up to 0.5s of connect_ex under flock(LOCK_EX), and the acquisition layer in #57687 calls janitor(), claim(), and refill() inside one locked block, re-polling every 0.25s. At the default pool size of 2 that's bounded at roughly a second, and the janitor limits how many stale members pile up, so this is a "worth saying out loud" item rather than a defect -- but spark.local.connect.pool.size is user-tunable, so the bound is too. Either note it in the docstring or restructure: select candidates under the lock, probe unlocked, then re-acquire to rename after re-checking the entry is still of kind server.
5. "Deterministic FIFO" rests on a wall clock. created is written with time.time() in #57687, so a backward clock adjustment (NTP step, suspend/resume) inverts the ordering. time.monotonic() isn't comparable across processes, so time.time() is the pragmatic choice -- I'd just say that in the docstring rather than claim determinism the data can't support. Ties are resolved by sorted() stability over sorted(os.listdir()), which is worth stating too since it's load-bearing and invisible.
Test coverage
Ten tests for this surface is solid, and the shape of the claim tests is right -- real subprocesses, real sockets, real files. The gaps below are all cheap:
- The
OverflowErrorcatch infrom_datais unreachable from any tested input. Every case ininvalid_recordsraisesPySparkValueErrororKeyError; the only way to hitOverflowErroris acreatedint too large to convert to float (10**400). Sincewrite_json/read_jsonusejsondefaults, such a value round-trips through a real state file, so the catch is what keeps a corrupt file from crashingclaim-- and nothing currently stops a refactor from deleting it. Relatedly,created = 2**100passes validation (finite as a float), and #57687's janitor computestime.time() - member.createdfor idle expiry, so a far-future member would never expire. An upper bound plus a test either way would settle it. float("inf")forcreatedis untested whilenanis. Sameisfinitecheck, different branch, and both are round-trippable:json.dumpsemitsInfinity/NaNandjson.loadreads them back, so this is a state file a real corruption can produce.claimoutside the lock isn't pinned.test_accessors_require_the_lockcoversuids(), andclaiminherits the assertion transitively viapaths_of_kind, but the docstring makes the lock a caller obligation. OneassertRaisesRegex(AssertionError, "context manager")onself._pool.claim("fp")would catch a future reordering that touches the directory before the first locked accessor.- No test asserts an already-
claimedmember is invisible toclaim. That's the core exclusion invariant. The kind filter makes it true, but the two-process test only demonstrates "claimed at most once" -- write aclaimed-<other-pid>-<uid>.jsonand assertclaimreturnsNone. - The
sorted()overseed_confis untested, and it's the entire reason the fingerprint is order-independent given that dicts preserve insertion order:pool_fingerprint(m, {"a": "1", "b": "2"}) == pool_fingerprint(m, {"b": "2", "a": "1"}). While there,{"k": 1}and{"k": "1"}collide throughstr()-- presumably intentional since confs serialize to a properties file, but worth an explicit assertion rather than an accident. sys.executableis in the identity list with no coverage.mock.patch.object(sys, "executable", "/other/python")is a one-liner, and it's the one identity component a packaging change could silently drop.test_concurrent_claimers_claim_one_member_oncedoesn't guarantee contention. If child A completes before B reachesflock, B sees noserverentry and printsNONE, satisfying the assertion without the two ever racing. It's a good regression test for "at most once"; a note saying so would keep a future reader from over-trusting it as a lock-contention test.
…print 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
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:
executables, PySpark and Spark paths, and
PYTHONPATH, using the full SHA-256 digest;PoolMemberfields;claims from separate processes.
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 ten focused tests at this layer, bringing the suite to 20 tests. They cover deterministic
Linux zombie detection, fingerprint identity and Python-interpreter precedence, Python and Spark
path compatibility, strict member-record validation, fingerprint-aware and FIFO claiming, a real
two-process claim race, unreachable members, and malformed, dead, or version-incompatible records.
All 20 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 alsopassed.
GitHub Actions run 31633722032
completed successfully: 19 checks passed and 5 inapplicable checks were skipped. The Apache-side
aggregate
Buildcheck also passed.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Fable 5) and OpenAI Codex (GPT-5)