Repository navigation
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces automatic retries for Managed Spark Sessions that fail to start during getOrCreate(). It adds a new SessionNotActiveError exception, which provides detailed error messages from Dataproc when a session fails or is deleted. A new configuration variable, MANAGED_SPARK_CONNECT_SESSION_START_RETRIES, allows users to specify the number of creation retries for generated session IDs. Additionally, the PR adds instrumentation for tracking session startup failures and includes comprehensive unit tests to verify the retry and error-handling behavior. I have no feedback to provide as there are no review comments.
An RPC on a session that Dataproc reports gone used to fail with a bare
RuntimeError("Session not active. Please create a new session"), dropping
the state the liveness check had just fetched. It now raises
SessionNotActiveError, which names the session and what Dataproc said:
"Managed Spark Session sc-123 is FAILED: <state_message>", "... is
TERMINATED", "... was deleted". The liveness cache keeps that detail with
a dead verdict, so RPCs that fast-fail on the cached verdict say the same.
SessionNotActiveError subclasses RuntimeError, as before, and
ManagedSparkConnectException, and its message keeps the old prefix and
suffix.
getOrCreate() connected to whatever session the CreateSession operation returned, without looking at its state or its Spark Connect endpoint. Reusing a session already checks both; a new one now does too: a session that is already FAILED, TERMINATING or TERMINATED raises SessionNotActiveError with what Dataproc reported, and one that is still CREATING or has no endpoint yet is polled until it has one. Polling for the endpoint also checked the endpoint before the state, so it returned a session that had failed after publishing its endpoint. It now checks for FAILED, TERMINATING (new) or TERMINATED first, and raises SessionNotActiveError with the state message instead of a bare ManagedSparkConnectException (which SessionNotActiveError subclasses).
…turns getOrCreate() returned a new session without any RPC to it outside notebooks, so a session that came up broken was only found out by the caller's first query: in CI, an 11-minute CreateSession, then a WebSocket handshake timeout and "Session not active" on the first createDataFrame(). getOrCreate() now makes the new session's first RPC itself: the Spark version, which the session-created message already fetched in notebooks and now gets from here. If it fails, or takes more than 60s (enough for a 30s WebSocket handshake timeout plus the liveness guard), Dataproc is asked about the session; one it reports gone is closed client-side and getOrCreate() raises SessionNotActiveError saying why. Anything else, including Dataproc being unreachable, returns the session as before. The check is recorded as a session.check_new span (first_rpc_ok, dead, timed_out). Tests that mock the Spark version now return one, since a version is what tells getOrCreate() that the new session is usable.
When a session getOrCreate() just created turns out dead before it could be used (Dataproc reports it failed or terminated in the CreateSession result, while waiting for its endpoint, or after its first RPC failed), getOrCreate() now creates another session in its place and returns that, up to MANAGED_SPARK_CONNECT_SESSION_START_RETRIES times (default 1, 0 to disable). Nothing has run on the dead session, so nothing is lost but time; the dead session itself is left alone, since Dataproc already has it failed or terminated. Only sessions with a generated ID are replaced: a session created with sessionId() belongs to the caller, so its failure still raises SessionNotActiveError. A replacement logs a WARNING, shows a status line in notebooks, and is recorded as a session.start_failed event (detail, attempt, retry).
dborowitz
force-pushed
the
dborowitz-session-start
branch
from
November 2, 2026 14:21
0349168 to
ae55192
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #217.
CI integration tests intermittently fail the same way on every OS: a new session takes ~10 minutes to create, the first RPC's WebSocket handshake times out, and the liveness guard raises a bare
RuntimeError("Session not active. Please create a new session"), because Dataproc by then reports the session not ACTIVE. The library returned a broken session fromgetOrCreate()and threw away what Dataproc said about it.This makes the library notice a new session that never came up, say why, and replace it when it can:
SessionNotActiveError(aRuntimeError, as before, and aManagedSparkConnectException) with what Dataproc said, e.g.Session not active: Managed Spark Session sc-123 is FAILED: <state_message>. Please create a new session.CreateSessionresult before connecting, as reuse already does: a session already FAILED/TERMINATING/TERMINATED raises; one without an endpoint yet is polled. The endpoint poll also checked the endpoint before the state, so it could return a session that had failed after publishing its endpoint; fixed.getOrCreate()returns, by making its first RPC (the Spark version, which notebooks already fetched for the session-created message). If it fails or takes over 60s, Dataproc is asked; a session it reports gone raisesSessionNotActiveErrorfromgetOrCreate()instead of from the caller's first query. Otherwise the session is returned as before. Recorded as asession.check_newspan.MANAGED_SPARK_CONNECT_SESSION_START_RETRIEStimes (default 1,0disables). Sessions created withsessionId()raise instead. Recorded as asession.start_failedevent; logs aWARNINGand shows a status line in notebooks.Not covered: a session Dataproc keeps reporting ACTIVE while its endpoint never accepts a WebSocket connection (seen once in the live run below). The check fails open there by design; the caller's first query still hangs on handshake timeouts.
Testing
tests/unit/test_session_start.pyand guard tests; full suite passes on pyspark-client 4.0.3.datagravity-demo,us-east4) with the internal resilience suite's newsession_start.*scenarios, which kill or cut off the first new session in a subprocess:dead_before_first_rpc(recreate): ok, replaced by a second session (2 sessions, onesession.start_failedwithretry=true, query ok).dead_before_first_rpc(no_retries),(custom_id):SessionNotActiveErrorfromgetOrCreate(),... is TERMINATING, no replacement.unreachable(until_ready): ok, first session returned after the 60s check (Dataproc reported it ACTIVE).terminated_mid_session(default),(liveness_ttl0):SessionNotActiveErrorwith... is TERMINATED/... is TERMINATING.fault.guard_transient_503(3 sizes),fault.ws_drop_while_idle,progress.driver_onlyon ws+guard and ws+noguard: unchanged (ok). The primary session's first-RPC check cost 1.7s of a 48sgetOrCreate().