Repository navigation
perf: eliminate per-RPC control-plane polling, cache auth/clients, and optimize proxy framing - #211
Merged
Merged
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 significant improvements to session liveness management and credential handling in the managed_spark_connect client. Key changes include the implementation of a liveness guard with TTL caching to prevent premature session termination, optimized credential caching, and refined socket handling for improved performance. The review feedback highlights a potential issue where disabling the fallback cache during error handling could lead to false-positive session terminations, suggesting that the fallback cache should remain enabled in on_error mode to allow for safer retry behavior.
dborowitz
force-pushed
the
dborowitz-optimization
branch
from
October 20, 2026 04:42
712aab3 to
2c55003
Compare
Contributor
Author
|
CI didn't run because it's configured to only run for PRs based on main; not surprising since the CI implementation predates stacked PRs as a feature. I'm folding a fix for that into #207. |
dborowitz
force-pushed
the
dborowitz-optimization
branch
from
October 22, 2026 20:59
2c55003 to
9c84762
Compare
dborowitz
force-pushed
the
dborowitz-optimization
branch
from
October 26, 2026 18:40
9c84762 to
280d9cc
Compare
dborowitz
force-pushed
the
dborowitz-optimization
branch
2 times, most recently
from
October 26, 2026 19:53
c91abcc to
4c6dcec
Compare
dborowitz
force-pushed
the
dborowitz-optimization
branch
from
October 27, 2026 19:32
4c6dcec to
d57d1d8
Compare
tim-u
reviewed
Oct 27, 2026
… endpoint Previously, every control-plane call and WebSocket tunnel connection fetched Application Default Credentials (`google.auth.default()`), refreshed OAuth2 tokens (`creds.refresh()`), and constructed a fresh Dataproc `SessionControllerClient` from scratch. In developer and workstation environments using user credentials, each uncached `google.auth.default()` call spawns `gcloud` subprocesses (~740-835 ms), each `SessionControllerClient()` construction without explicit credentials spawns `gcloud` and mTLS helper subprocesses while opening an unclosed gRPC TLS channel (~750-1,680 ms), and each unconditional `creds.refresh()` adds an ~80-90 ms HTTP round-trip. Even without per-RPC liveness polling, a single session lifecycle still loads credentials and instantiates `SessionControllerClient` multiple times across `session.py`, `client/core.py`, and `client/proxy.py` when creating or attaching to a session, opening the WebSocket proxy on the first RPC, and polling termination in `spark.stop()`. To share credentials across the WebSocket proxy, direct gRPC channels, and Dataproc control-plane clients without circular imports, credential caching lives in `client/proxy.py`. `get_default_credentials()` caches `(creds, project)` under a re-entrant lock (`_AUTH_LOCK`) keyed by OAuth scopes and auth environment variables (`GOOGLE_APPLICATION_CREDENTIALS`, `GOOGLE_CLOUD_PROJECT`, `HOME`). `get_refreshed_credentials()` reuses those cached credentials and skips `creds.refresh()` under `_AUTH_LOCK` while `creds.valid` is `True` (which enforces the `google-auth` built-in 5-minute clock-skew margin against `creds.expiry`) and `creds.token` is populated, clearing `_CREDENTIALS_CACHE` on refresh failure so broken credentials are never pinned. In `client/core.py`, `ManagedSparkChannelBuilder._direct_channel` also reuses `get_default_credentials()` and builds composite SSL and `AuthMetadataPlugin` call credentials with a default port of 443. In `session.py`, `_create_session_controller_client()` caches `SessionControllerClient` instances in `_CLIENT_CACHE` keyed by a normalized tuple of endpoint and auth options (`api_endpoint`, `quota_project_id`, `universe_domain`, `api_audience`, `credentials_file`, `api_key`, `scopes`) and environment variables. When standard default credentials are used, it passes the cached `creds` from `get_default_credentials()` into `SessionControllerClient` so instantiation skips redundant `google.auth.default()` subprocesses and the gRPC auth metadata plugin refreshes expiring tokens in-flight on the open channel; when custom credential options or client certificate callbacks are supplied in `client_options`, it avoids passing default credentials so `google-api-core` does not raise `DuplicateCredentialArgs`. Both caches automatically bypass caching when `google.auth.default` or `SessionControllerClient` is monkeypatched in unit tests, and `clear_dataproc_caches()` provides explicit cache reset. Simpler alternatives—such as caching the client on the `Builder` instance or threading a client parameter through internal helpers—are insufficient because `SparkSession.builder` constructs a new `Builder` on every property access, while session reuse checks, background WebSocket proxy reconnect threads, liveness checks, and `terminate_s8s_session()` in `spark.stop()` all run outside the original `Builder` instance. Process-wide caching keyed by effective endpoint and auth configuration eliminates duplicate subprocesses and gRPC channel leaks across all lifecycle entry points without changing public APIs. Benchmark summary: Reduces warm credential lookup and client acquisition from ~745-827 ms (3-6 subprocesses) to under 2 us (0 subprocesses), cutting first Spark Connect RPC latency (including WebSocket handshake auth) by 63% (from 4.04 s to 1.51 s, 2.7x faster) and eliminating ~705-813 ms of client initialization overhead during `spark.stop()`.
…aseExecute Previously, `ProxiedChannel` wrapped every Spark Connect gRPC method in `_GuardedMultiCallable`, which executed a synchronous Dataproc control-plane `GetSession` RPC before every unary and streaming call. Because a single DataFrame operation issues 3-7 foreground Spark Connect RPCs plus 8-100+ background `ReleaseExecute` cleanup RPCs (which account for over 90% of all gRPC calls), this pre-RPC guard added ~100-350 ms of control-plane RPC latency (and over 1,200 ms prior to client caching) to every ~50 ms Spark operation, while concurrent background `ReleaseExecute` guard threads contended on the Python GIL and slowed Arrow result decoding by more than 2x. In addition, `get_active_s8s_session_response()` caught all exceptions and returned `None`, causing transient control-plane 503/504/429 blips to abort healthy Spark sessions with `RuntimeError(Session not active)`. In `client/core.py`, `ProxiedChannel._wrap_method()` now exempts background `/ReleaseExecute` and `/ReleaseSession` cleanup calls from guard wrapping altogether. For foreground RPCs, `_GuardedMultiCallable` and `_GuardedStream` coordinate with the guard callback via a `_GUARD_MODE` `ContextVar` (`pre_rpc` vs. `on_error`). On the happy path, `_run_pre_guard()` executes in `pre_rpc` mode with zero `GetSession` RPCs and dispatches the Spark Connect call directly, wrapping streaming responses (`ExecutePlan`) in `_GuardedStream` so mid-stream failures are also intercepted. Only when a call raises an exception that `_should_check_liveness_on_error()` classifies as a possible session termination—skipping client cancellations, invalid argument errors, and rich gRPC errors carrying `grpc-status-details-bin` Spark SQL/execution exceptions from a live server—does `_run_error_guard()` invoke the guard in `on_error` mode. In `session.py`, `make_is_active_callback()` tracks session state across `_DEAD_SESSIONS`, `_LAST_KNOWN_ACTIVE`, and `_LIVENESS_CACHE`. When an `on_error` check confirms via `GetSession` that a session is no longer active, the session is recorded in `_DEAD_SESSIONS` and `ProxiedChannel` raises `RuntimeError(Session not active. Please create a new session)` before the PySpark `RetryPolicy` can enter a ~60 s retry backoff loop; all subsequent RPCs on that session fast-fail in `pre_rpc` mode without network calls, whereas if the session is still `ACTIVE`, the original `grpc.RpcError` is re-raised unmodified. To harden control-plane checks, `get_active_s8s_session_response()` falls back to `_LAST_KNOWN_ACTIVE` for up to 60 s on transient `503`/`504`/`429`/`500` control-plane errors instead of falsely marking the session dead, and supports an optional `MANAGED_SPARK_CONNECT_LIVENESS_TTL_SECONDS` environment variable if periodic pre-RPC TTL polling is ever desired. Checking liveness lazily on RPC error was chosen over pre-RPC TTL caching, background polling threads, or removing the guard entirely. Pre-RPC TTL caching and background polling still require on-error checks when a session terminates mid-TTL window or mid-stream, while introducing periodic tail-latency spikes on TTL expiry or background thread lifecycle complexity. Conversely, removing the guard altogether loses actionable session-termination diagnostics because a terminated session drops the WebSocket proxy with a generic `UNAVAILABLE` error that causes PySpark to retry repeatedly before surfacing an opaque transport failure. Benchmark summary: Eliminates 100% of control-plane `GetSession` RPCs on the happy path, reducing end-to-end latency by 80-90% (5.0-10.0x faster) on interactive schema, config, and SQL operations, 68-90% (2.8-10.3x faster) on DataFrame queries and UDFs, 71% (3.5x faster) on a 100-iteration query loop, and 81-85% (5.1-6.6x faster) when propagating Spark SQL analysis and runtime errors.
…x partial send In the default proxy mode, `ManagedSparkSessionProxy` (`client/proxy.py`) bridges local gRPC/HTTP2 traffic over a WebSocket tunnel, placing it in the data path of every Spark Connect request and response. Previously, four compounding issues in this bridge severely throttled throughput, inflated small-RPC latency, and risked stream corruption: `forward_bytes` read from the socket in 1,024-byte chunks (`recv(1024)`), turning a 4 MiB upload into ~4,096 separate WebSocket frames that stalled behind synchronous round-trips in the remote proxy; `websockets.sync.client` left `permessage-deflate` enabled by default, burning single-threaded Python zlib CPU on already-compressed Arrow and protobuf payloads; accepted loopback sockets left the Nagle algorithm enabled, triggering ~44 ms delayed-ACK stalls on multi-chunk messages; and `forward_bytes` called `to_sock.send(bs)` only once without looping, which can silently truncate HTTP/2 frames whenever the kernel TCP send buffer partially accepts a large inbound WebSocket frame. To fix these bottlenecks and the partial-write bug in `client/proxy.py`, `RECV_BUFFER_SIZE` is increased to `65536` (matching the default 64 KiB HTTP/2 window), and `connect_tcp_bridge` passes `compression=None` and `max_size=MAX_WS_MESSAGE_SIZE` (`64 * 1024 * 1024`, a 64 MiB cap that supports large Arrow batches without unbounded memory growth). When accepting local connections in `forward_connection` (and shutting down the listener cleanly in `ManagedSparkSessionProxy.stop`), the proxy sets `TCP_NODELAY` on each loopback socket. Finally, `bridged_socket.send` now returns `len(msg_bytes)` and `forward_bytes` loops over remaining byte slices on partial `to_sock.send()` calls—resetting the timeout retry counter whenever forward progress is made, raising `ConnectionError` if `send()` returns `0` on a closed peer, and closing both sockets on unrecoverable errors. Tuning the synchronous socket bridge and fixing its partial-write loop was chosen over replacing `websockets.sync` with an `asyncio` proxy or switching entirely to direct gRPC. While direct gRPC is still blocked in environments where upstream HTTP/2 proxies strip trailing gRPC error details, these targeted fixes eliminate almost all local framing and compression overhead with minimal invasive changes while preserving the existing threaded proxy architecture. Benchmark summary: - Local loopback WebSocket proxy throughput increases 8.0x-18.8x on bulk 35.5 MB transfers (reducing process CPU time by 87-95%), and 8 KiB ping-pong RPC latency drops by 98% (from 44.0 ms to 0.70 ms). - End-to-end live Dataproc Spark Connect latency improves by 17-30% on small interactive SQL/catalog/schema RPCs and pandas UDFs (proxy-only), and by 43-86% (1.8x-7.2x faster) when downloading Arrow result batches into pandas alongside the liveness guard fix.
…ctive session Attaching to an already-running Dataproc session via `ManagedSparkSession.builder.sessionId(...).getOrCreate()` previously made three sequential `GetSession` control-plane RPCs (~250-300 ms each) and constructed two separate `SessionControllerClient` instances before sending the first Spark Connect RPC, even when the session was already `ACTIVE` with its Spark Connect endpoint populated on the very first call. Specifically, `_handle_custom_session_id` fetched the session to validate its state and discarded the response; `_get_exiting_active_session` immediately instantiated a new client to fetch the same session a second time; and `_wait_for_session_available` unconditionally polled `GetSession` a third time to read the endpoint map. To eliminate the redundant round-trips in `session.py`, `_handle_custom_session_id` now stashes the `Session` proto from its initial `GetSession` call in `self._custom_session_response` whenever `state == Session.State.ACTIVE` (clearing it on entry and in a `finally` block in `getOrCreate()` so stale responses never leak across builds) and calls `mark_s8s_session_active` so the liveness tracker knows the session is live without an extra check. When switching between distinct session IDs or UUIDs, `_handle_custom_session_id` also cleanly detaches any previous local client wrapper (`stop(terminate=False)`). Downstream, `_get_exiting_active_session` consumes `self._custom_session_response` instead of issuing a second `GetSession` RPC, and checks whether `Spark Connect Server` is already present in `session_response.runtime_info.endpoints`—calling `__create_spark_connect_session_from_s8s` directly and only falling back to the `_wait_for_session_available` polling loop if the endpoint map is not yet populated. In addition, `Builder.__create` now reuses `self.session_controller_client` instead of constructing a duplicate client when creating a new session. Passing the already-fetched `Session` response across the existing two-stage validation and attachment methods was chosen over refactoring the entire session builder call graph. This preserves all existing validation, telemetry spans, and fallback polling behavior when a target session is still `CREATING` or waiting for its endpoint to register, while cutting steady-state attachment down to a single `GetSession` RPC. Benchmark summary: - Cuts `GetSession` control-plane RPCs when attaching to an active session from 3 to 1 (-67%) and `SessionControllerClient` creations from 2 to 1 (-50%). - Reduces in-process `getOrCreate()` control-plane attachment latency by 65% (from 262 ms to 91 ms at 85 ms RTT) and live end-to-end existing-session attachment time by 61-67% (2.5x-3.1x faster).
…e pbar on event Session lifecycle operations previously incurred unnecessary dead wait time from three fixed-sleep and polling patterns. First, when stopping a session, `terminate_s8s_session()` polled `get_session()` in a `while` loop until the session entered `TERMINATING`, `TERMINATED`, or `FAILED`, but executed an unconditional `time.sleep(1)` at the bottom of the loop body before re-evaluating the loop condition. Because Dataproc transitions sessions to `TERMINATING` almost immediately during `terminate_session()`, the first `get_session()` call already observed `TERMINATING`, yet every `stop()` call still slept a full 1.0s before returning (and had no upper bound if a session never transitioned). Second, during session creation, the background progress-bar thread (`create_session_pbar`) used uninterruptible `time.sleep(1)` calls, blocking `create_session_pbar_thread.join()` for up to 1.0s (500ms median) after the session was already ready, while `operation.result()` polled the creation operation at a fixed 5.0s interval (`initial=5.0`, `multiplier=1.0`), adding up to 5.0s of polling quantization delay. In `terminate_s8s_session()`, the loop now checks `state not in terminal_states` with a `max_polls = 180` safety bound and only calls `time.sleep(1)` after `get_session()` if the session has not yet reached `TERMINATING`, `TERMINATED`, or `FAILED` and further poll iterations remain. In `Builder.__create()`, `create_session_pbar` replaces `time.sleep(1)` with `stop_create_session_pbar_event.wait(1)` (and `stop_create_session_pbar_event.wait()` on the final progress tick) so that setting `stop_create_session_pbar_event` upon session creation completion wakes and joins the progress-bar thread immediately. Finally, the `retry.Retry` policy in `Builder.__create()` is updated to start at `initial=1.0` with `multiplier=1.5` up to `maximum=5.0`. Checking `get_session()` once immediately after `terminate_session()` and sleeping only when another poll is actually required was chosen over removing the post-termination check altogether because it preserves the guarantee that the session has transitioned out of `ACTIVE` before `stop()` returns, while eliminating the 1.0s post-check penalty in the common case and bounding worst-case polling at 180 attempts. Similarly, replacing fixed 5.0s polling with exponential backoff (`1.0s` to `5.0s`) and event-driven progress-bar wakeup detects newly ready sessions several seconds earlier without flooding the Dataproc control plane during longer cluster startups. Benchmark summary: - Stopping a session end-to-end is 46-55% faster (1.8-2.2x speedup, saving ~1.0s per `stop()` call), with the termination state-check step dropping by 88% (from ~1,112ms to ~137ms). - Progress-bar thread teardown at the end of session creation drops from ~500ms median (up to 1,000ms) to ~0.5ms, and initial readiness polling starts at 1.0s instead of 5.0s.
…ng.requirements Importing `google.cloud.managed_spark_connect` in non-interactive CLI scripts, batch jobs, or services previously triggered avoidable module loading when optional dependencies were installed in the environment. At package import time, `__init__.py` unconditionally called `_in_ipython()` (which executed `from IPython import get_ipython`) and imported `_ipython._init_extras()` whenever `traitlets` was installed, pulling in `IPython` and dozens of its transitive submodules (~100-200ms of isolated import overhead) even when running in a standard Python script outside any notebook or IPython shell. In addition, `pypi_artifacts.py` imported `packaging.requirements.Requirement` eagerly at module top level (loaded transitively via `session.py`), pulling in the `packaging` parser and tokenizer submodules even when PyPI package artifacts were never used. To avoid loading interactive notebook machinery outside an active IPython session, `_in_ipython()` in `__init__.py` and `_init_extras()` in `_ipython.py` now look up `sys.modules.get(IPython)` before calling `get_ipython()`, and `__init__.py` gates importing `_ipython` and calling `_init_extras()` on `_in_ipython() and importlib.util.find_spec(traitlets) is not None`. Since an active Jupyter or IPython kernel always has `IPython` present in `sys.modules` before user code imports the library, `sys.modules.get(IPython)` accurately detects active interactive shells without triggering a cold import of `IPython`. In `pypi_artifacts.py`, the top-level `packaging.requirements.Requirement` import is placed under `if TYPE_CHECKING:` for static type checkers and deferred to runtime inside `PyPiArtifacts.__try_parsing_package()`. Inspecting `sys.modules.get(IPython)` before invoking `get_ipython()` was chosen over removing automatic interactive initialization or relying solely on `importlib.util.find_spec(traitlets)` because virtual environments frequently have `IPython` and `traitlets` installed even when executing non-interactive batch scripts. Checking `sys.modules` completely eliminates the `IPython` and `_ipython` import penalty for standard Python scripts while still automatically registering magics and display integrations whenever the package is imported inside an active IPython or Jupyter notebook session. Benchmark summary: - Eliminates 86 transitive module imports on non-notebook startup, reducing isolated interactive-check overhead by 50% (~207ms to ~103ms) and `pypi_artifacts` module import time by 97% (~17.1ms to ~0.4ms).
dborowitz
force-pushed
the
dborowitz-optimization
branch
from
October 27, 2026 23:51
d57d1d8 to
0d5298c
Compare
tim-u
approved these changes
Oct 28, 2026
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.
Note
Stacked on #210 (
dborowitz-timing).Summary
Eliminates client-side control-plane, authentication, proxy framing, session attachment, polling, and import-time bottlenecks across
google-cloud-spark-connect.Prior to this series, even simple interactive Spark Connect operations suffered from severe client-side overhead:
ReleaseExecutecleanup RPC) synchronously calledGetSessionon the Dataproc control plane and constructed a freshSessionControllerClientfrom scratch.gcloud/ mTLS subprocesses (google.auth.default()) and unconditionally refreshed OAuth2 tokens (creds.refresh()).ManagedSparkSessionProxy) read in 1 KiB chunks (recv(1024)), compressed already-compressed Arrow/protobuf frames with single-threaded Pythonzlib(permessage-deflate), left Nagle's algorithm enabled on loopback sockets, and did not loop on partialsocket.send()writes.ManagedSparkSession.builder.sessionId(...).getOrCreate()) issued 3 sequentialGetSessionRPCs and built 2 clients even when the first response already showed the session wasACTIVEwith its endpoint populated.spark.stop()slept an unconditional 1.0s after confirming the session had already enteredTERMINATING, and session creation blocked up to 1.0s waiting for the progress-bar thread'stime.sleep(1)to wake.google.cloud.managed_spark_connectin non-interactive scripts eagerly importedIPython(~100–200ms and 86 transitive modules) whenever optional interactive dependencies were installed in the environment.End-to-End Benchmark Highlights
Benchmarked using a custom suite based on the instrumentation harness from #210, with top issues identified and fixed by my agent. I do intend to make the benchmark suite public, but it's another few thousand lines of AI-generated code that's not in the trivial path, and needs at least cursory human review by me before sending.
df.columns,spark.conf.get,SELECT 1)range, joins, aggregations, Python/pandas UDFs)pandas(toPandas())ACTIVEsession (getOrCreate())GetSessionRPCsGetSessionRPCspark.stop())Key Changes (by Commit)
1.
perf(auth): cache default credentials and SessionControllerClient per endpointgoogle.auth.default()(spawninggcloudsubprocesses, ~740–835 ms),creds.refresh()(~80–90 ms HTTP round-trip), andSessionControllerClient()without explicit credentials (spawninggcloud/mTLS helpers and opening an unclosed gRPC TLS channel, ~750–1,680 ms).get_default_credentials()inclient/proxy.pycaches(creds, project)under a re-entrant lock (_AUTH_LOCK) keyed by OAuth scopes and auth environment variables (GOOGLE_APPLICATION_CREDENTIALS,GOOGLE_CLOUD_PROJECT,HOME).get_refreshed_credentials()reuses cached credentials and only callscreds.refresh()whennot creds.validornot creds.token(honoringgoogle-auth's 5-minute clock-skew margin), clearing the cache on refresh error._create_session_controller_client()insession.pycachesSessionControllerClientinstances in_CLIENT_CACHEkeyed by normalized endpoint/auth options and passes cached default credentials so client initialization skips redundant subprocesses.google.auth.defaultorSessionControllerClientis monkeypatched in unit tests, and providesclear_dataproc_caches()for explicit resets.2.
perf(guard): check session liveness lazily on RPC error and skip ReleaseExecuteProxiedChannelexecuted a synchronous DataprocGetSessionRPC before every unary and streaming Spark Connect call—including backgroundReleaseExecutecleanup RPCs (which represent >90% of all gRPC calls)—adding ~100–350 ms of control-plane latency (and >1,200 ms before client caching) to every ~50 ms Spark operation and contending on the GIL during Arrow decoding. Furthermore, transient control-plane503/504/429errors causedget_active_s8s_session_response()to returnNoneand falsely abort healthy sessions./ReleaseExecuteand/ReleaseSessioncleanup RPCs from guard wrapping inProxiedChannel._wrap_method()._GUARD_MODEContextVar(pre_rpcvs.on_error) between_GuardedMultiCallable/_GuardedStream(client/core.py) andmake_is_active_callback()(session.py). On the happy path,pre_rpcdispatches the Spark Connect RPC with zeroGetSessioncalls (unless the session is already known dead in_DEAD_SESSIONSor an optionalMANAGED_SPARK_CONNECT_LIVENESS_TTL_SECONDSis configured)._should_check_liveness_on_error()classifies as a possible session termination (skipping client cancellations, invalid arguments, and rich gRPC errors carryinggrpc-status-details-binSpark exceptions from a live server),_run_error_guard()checksGetSessioninon_errormode. If the session is no longer active, it records the session in_DEAD_SESSIONSand raisesRuntimeError("Session not active. Please create a new session")before PySpark'sRetryPolicycan enter a ~60s retry loop; if the session is stillACTIVE, the originalgrpc.RpcErroris re-raised unmodified.get_active_s8s_session_response()to fall back to_LAST_KNOWN_ACTIVEfor up to 60s on transient503/504/429/500control-plane errors.3.
perf(proxy): use 64KiB recv, disable deflate, set TCP_NODELAY, and fix partial sendManagedSparkSessionProxy(client/proxy.py) read from sockets withrecv(1024)(fragmenting a 4 MiB payload into ~4,096 WebSocket frames), left WebSocketpermessage-deflateenabled (burning single-threaded PythonzlibCPU on Arrow/protobuf data), left Nagle's algorithm enabled on accepted loopback sockets (~44 ms delayed-ACK stalls), and calledto_sock.send(bs)only once without looping (risking silent HTTP/2 frame truncation on partial kernel socket writes).RECV_BUFFER_SIZEto65536(matching the 64 KiB HTTP/2 window).compression=Noneandmax_size=MAX_WS_MESSAGE_SIZE(64 MiB) inconnect_tcp_bridge.socket.TCP_NODELAYon accepted loopback connections inforward_connectionand shuts down the listener socket cleanly inManagedSparkSessionProxy.stop().forward_byteson partialto_sock.send()calls, resetting the timeout counter on forward progress and raisingConnectionErrorifsend()returns0.4.
perf(session): avoid redundant GetSession RPCs when attaching to an active sessionManagedSparkSession.builder.sessionId(...).getOrCreate()made 3 sequentialGetSessionRPCs (_handle_custom_session_id->_get_exiting_active_session->_wait_for_session_available) and constructed 2SessionControllerClientinstances even when the session was alreadyACTIVEwith its Spark Connect endpoint present on the first call.Sessionproto from_handle_custom_session_idinself._custom_session_responsewhenstate == Session.State.ACTIVE(cleaned up on entry and infinally:ofgetOrCreate()), marks the session active in the liveness tracker, and cleanly detaches stale local client state (stop(terminate=False)) when switching session IDs or UUIDs.self._custom_session_responsein_get_exiting_active_sessionand calls__create_spark_connect_session_from_s8sdirectly when"Spark Connect Server"is already present insession_response.runtime_info.endpoints, falling back to_wait_for_session_availableonly if the endpoint is not yet populated.self.session_controller_clientinBuilder.__create.5.
perf(polling): skip post-check sleep in terminate_s8s_session and wake pbar on eventterminate_s8s_session()slept an unconditional1.0sat the end of itswhileloop body even though the firstget_session()afterterminate_session()already observedTERMINATING;create_session_pbarused uninterruptibletime.sleep(1)calls that blockedjoin()for up to 1.0s after session creation finished; andoperation.result()polled at a flat 5.0s interval.terminate_s8s_session()whenstatehas not yet reachedTERMINATING,TERMINATED, orFAILED(with amax_polls = 180safety bound).time.sleep(1)increate_session_pbarwithstop_create_session_pbar_event.wait(1)so setting the stop event wakes the progress-bar thread immediately.retry.Retrypolicy inBuilder.__createto start atinitial=1.0withmultiplier=1.5up tomaximum=5.0.6.
perf(import): skip IPython import outside notebooks and defer packaging.requirementsgoogle.cloud.managed_spark_connectunconditionally executedfrom IPython import get_ipythonand imported_ipythonwhenevertraitletswas installed, loadingIPythonand 86 transitive modules (~100–200ms) even in non-interactive CLI scripts and batch jobs;pypi_artifacts.pyalso eagerly importedpackaging.requirements.Requirementat module import time.sys.modules.get("IPython")in_in_ipython()(__init__.py) and_init_extras()(_ipython.py) before invokingget_ipython(), avoiding cold-importingIPythonoutside active interactive kernels while preserving automatic magic/display registration inside Jupyter and IPython sessions.from packaging.requirements import Requirementinpypi_artifacts.pytoPyPiArtifacts.__try_parsing_package().Testing
tests/unit/test_dataproc.py,tests/unit/test_guard.py,tests/unit/test_proxy.py,tests/unit/test_session.py,tests/unit/test_polling.py,tests/unit/test_init.py, andtests/unit/test_core.py.