Skip to content

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

Open
ericm-db wants to merge 6 commits into
apache:masterfrom
ericm-db:local-connect-pool-claiming
Open

[SPARK-58021][CONNECT] Add local server pool member claiming#57907
ericm-db wants to merge 6 commits into
apache:masterfrom
ericm-db:local-connect-pool-claiming

Conversation

@ericm-db

@ericm-db ericm-db commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

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
  • deterministic FIFO claiming through an atomic state-file rename, including safe concurrent
    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.

python/run-tests --testnames pyspark.sql.tests.connect.test_connect_local_server_pool

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 also
passed.

GitHub Actions run 31633722032
completed successfully: 19 checks passed and 5 inapplicable checks were skipped. The Apache-side
aggregate Build check also passed.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Fable 5) and OpenAI Codex (GPT-5)

dtenedor pushed a commit that referenced this pull request Aug 12, 2026
### 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>
dtenedor pushed a commit that referenced this pull request Aug 12, 2026
### 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>
@ericm-db
ericm-db force-pushed the local-connect-pool-claiming branch from 06648f8 to 31d64ca Compare August 12, 2026 19:21

@dtenedor dtenedor left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

…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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants