[SPARK-58021][CONNECT] Add local server pool member claiming

[ericm-db]

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.

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 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)

[dtenedor]

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 True

It 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 OverflowError catch in from_data is unreachable from any tested input. Every case in invalid_records raises PySparkValueError or KeyError; the only way to hit OverflowError is a created int too large to convert to float (10**400). Since write_json / read_json use json defaults, such a value round-trips through a real state file, so the catch is what keeps a corrupt file from crashing claim -- and nothing currently stops a refactor from deleting it. Relatedly, created = 2**100 passes validation (finite as a float), and #57687's janitor computes time.time() - member.created for idle expiry, so a far-future member would never expire. An upper bound plus a test either way would settle it.
• float("inf") for created is untested while nan is. Same isfinite check, different branch, and both are round-trippable: json.dumps emits Infinity / NaN and json.load reads them back, so this is a state file a real corruption can produce.
• claim outside the lock isn't pinned. test_accessors_require_the_lock covers uids(), and claim inherits the assertion transitively via paths_of_kind, but the docstring makes the lock a caller obligation. One assertRaisesRegex(AssertionError, "context manager") on self._pool.claim("fp") would catch a future reordering that touches the directory before the first locked accessor.
• No test asserts an already-claimed member is invisible to claim. That's the core exclusion invariant. The kind filter makes it true, but the two-process test only demonstrates "claimed at most once" -- write a claimed-<other-pid>-<uid>.json and assert claim returns None.
• The sorted() over seed_conf is 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 through str() -- presumably intentional since confs serialize to a properties file, but worth an explicit assertion rather than an accident.
• sys.executable is 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_once doesn't guarantee contention. If child A completes before B reaches flock, B sees no server entry and prints NONE, 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.

[ericm-db]

Thanks, @dtenedor - addressed all points in 5e92368f37f: shared probes, completed the fingerprint, documented locking/FIFO semantics, tightened timestamp validation, and added the requested tests. CI is green.

[dtenedor]

Thanks -- re-reviewed at 5e92368. All nine items from the last round are addressed, and I checked each
against the code rather than the commit message. Consolidation is the part I'd call done well: there is now
exactly one _pid_alive and one _port_open in local_server.py, the hardening (pid <= 0, OverflowError,
Linux zombies, the connect_ex guard) reaches the reuse path as well as the pool path, and is_usable() and
is_listening() both delegate. The fingerprint now carries both interpreter resolutions, and I confirmed the
two precedences the comment cites: SparkConnectPlanner.pythonExec is
sys.env.getOrElse("PYSPARK_PYTHON", sys.env.getOrElse("PYSPARK_DRIVER_PYTHON", "python3")) while
PythonUtils.defaultPythonExec is the reverse, so udf_python / data_source_python mirror them correctly and
the test that used to pin the gap as intended behavior now asserts assertNotEqual. The _JVM_ENV_VARS tuple
plus the softened "curated set, not an exhaustive one" docstring resolve the completeness promise, and the
created bound closes the far-future hole with the OverflowError guard now exercised by a real 10**400
record. Test count went from 10 to 25.

One substantive follow-up on the created bound, then three nits.

The _MAX_CREATED rejection moves the far-future leak from the claim path to the reap path. Rejecting the
whole record is right for claim: a corrupt member is no longer claimable, which was the important half.
But from_data returning None discards the pid, which in that record is usually perfectly readable, and
#57687's _reap_server routes exactly that case into self._retire(path, member.pid if member is not None else -1). So with created = 2**100, the behavior goes from "member never expires but stays tracked and
reachable" to "state file is retired promptly while its JVM keeps running untracked" -- _signal(-1, ...) is a
no-op (correctly guarded), and the next _reap_retired pass sees pid == -1, treats it as gone, and removes
the retired marker and the member directory. The server is then unreachable to every rule. Two ways out, both
cheap: keep the record parseable and treat an out-of-range created as age infinity (immediately idle-expired),
so the reaper still SIGTERMs the recorded pid; or have _reap_server fall back to the spark-daemon pid file in
member-<uid>/ for unparseable records, which _kill_recorded_daemon already reads on the pending path. This
is a #57687 change either way -- flagging it here because this PR is what starts routing records with a good
pid down that branch.

Nits:

• _await_ready still inlines the probe that _port_open now owns -- same AF_INET, same 0.5s, same
  connect_ex == 0, three lines above the version in the same file. It is the last copy, so
  _port_open("localhost", port) finishes the consolidation, and it also picks up the OSError guard: today a
  transient socket error there propagates out and aborts the launch instead of retrying until the deadline.

• The POSIX guard for _pid_alive is at a call site rather than in the helper. is_reusable still checks
  os.name == "posix" before calling it; PoolMember.is_usable() calls it unguarded. That asymmetry matters
  more than it looks because on Windows os.kill(pid, sig) for any sig other than CTRL_C_EVENT /
  CTRL_BREAK_EVENT "will cause the process to be unconditionally killed by the TerminateProcess API" -- so
  the probe is destructive there, not just wrong. It is unreachable today (the pool needs fcntl, and the suite
  is skipUnless(os.name == "posix")), but now that the helper is shared, the guard belongs inside it, which
  also lets is_reusable drop its own copy.

• test_fingerprint_includes_jvm_env covers 2 of the 7 _JVM_ENV_VARS. A subTest loop over the tuple is
  the same length as the two hand-written cases and fails if a variable is ever dropped from it, which is the
  regression the test exists to catch.

• Optional: PATH is the notable remaining omission from the fingerprint. bin/spark-class uses
  ${JAVA_HOME}/bin/java when JAVA_HOME is set and falls back to command -v java otherwise, so two runs
  with different JDKs first on PATH and no JAVA_HOME share a member. The identity already depends on PATH
  indirectly through shutil.which for the interpreters. The curated-set caveat covers this, so either adding
  it or naming it as a known exclusion in the _JVM_ENV_VARS comment is fine.

Two things outside the code:

• The PR description is now stale in three places. It says "Added ten focused tests at this layer, bringing
  the suite to 20 tests" (it is 15 and 25), it still advertises "deterministic FIFO claiming" where the code now
  says "well defined but only approximately FIFO, not guaranteed", and the linked Actions run predates
  5e92368. Worth fixing before merge since the description becomes the commit message.

• 5e92368's trailer is Co-authored-by: Isaac with no <email>, so it is malformed -- GitHub will not
  attribute it and it will land in the squashed commit as-is.

None of this blocks; the claiming protocol and the fingerprint contract both look right to me now.

[ericm-db]

@dtenedor, addressed in 36c787ab5d6:

1. _await_ready now uses the shared _port_open (the last inlined probe) and so gains its OSError/UnicodeError retry.
2. The POSIX guard moved into _pid_alive itself; is_reusable drops its own and shares the helper with PoolMember.is_usable.
3. test_fingerprint_includes_jvm_env now covers all seven _JVM_ENV_VARS via an independent expected list (the earlier loop over the tuple was circular).
4. PATH documented as a deliberate omission in the _JVM_ENV_VARS comment.
5. PR description fixed: 15 added / 25 total, "well-defined, approximately-FIFO", refreshed Actions link.
6. Co-author trailer dropped — sole author now.

One deferral to confirm: the _MAX_CREATED reap-leak belongs in the reaper (_reap_server / _kill_recorded_daemon), which is in #57687, not this branch, so I'll fix it there. Thanks for the thorough review!

[dtenedor]

LGTM, merging to master and 4.x

[dtenedor]

Merge Summary:

• merged into master ee11a92
• merged into branch-4.x 30f8931

Posted by merge_spark_pr.py