Repository navigation
feat(instrumentation): opt-in instrumentation of sessions, RPCs, proxy and client - #210
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 opt-in instrumentation for Managed Spark Connect sessions, enabling the recording of timed spans for session lifecycle events, gRPC calls, proxy connections, and client-side conversions. The feedback identifies three key areas for improvement: a potential race condition in _active_recorder() when accessing the shared _registered list without a lock, a performance overhead in the is_active liveness check caused by repeatedly instantiating a new SessionControllerClient, and a missing timeout in the terminate_s8s_session polling loop that could lead to indefinite hangs during session termination.
dborowitz
force-pushed
the
dborowitz-timing
branch
from
October 20, 2026 04:42
714ca39 to
ced52cd
Compare
dborowitz
force-pushed
the
dborowitz-timing
branch
from
October 22, 2026 20:59
ced52cd to
c1aa0a8
Compare
dborowitz
force-pushed
the
dborowitz-timing
branch
from
October 26, 2026 18:40
c1aa0a8 to
9ed312b
Compare
tim-u
reviewed
Oct 26, 2026
dborowitz
force-pushed
the
dborowitz-timing
branch
2 times, most recently
from
October 26, 2026 19:53
b3258e5 to
78b2072
Compare
tim-u
approved these changes
Oct 27, 2026
ProxiedChannel._wrap_method replaced each gRPC multi-callable with a plain function, dropping the with_call() and future() entry points. Anything that layers grpc.intercept_channel() on top of the channel (upstream ChannelBuilder.add_interceptor, tracing/metrics interceptors) calls with_call() for unary-unary and stream-unary RPCs and failed with an AttributeError. Wrap the multi-callable in a small object that runs the liveness guard before __call__, with_call and future alike. subscribe()/unsubscribe() return None and are now forwarded instead of being wrapped.
ManagedSparkChannelBuilder inherits add_interceptor() from the upstream ChannelBuilder, but toChannel() returned the ProxiedChannel directly, so registered interceptors were silently ignored (the inner channel is built by a fresh DefaultChannelBuilder with an empty interceptor list). Apply the builder's interceptors around the ProxiedChannel, matching what ChannelBuilder._insecure_channel/_secure_channel do upstream.
Add google.cloud.managed_spark_connect.instrumentation: a small, dependency-free tracing layer that records timed spans (plus events, counters and one meta record) to pluggable sinks (JSONL file, in-memory, logging, callback). Off by default: without an Instrumentation the library goes through a shared no-op recorder, installs no gRPC interceptor and leaves the pyspark client untouched. Opt in with ManagedSparkSession.builder.instrumentation(inst), or via MANAGED_SPARK_CONNECT_INSTRUMENTATION_FILE / _LOG. The session exposes `session.instrumentation` and closes it on stop(). What is recorded: - session lifecycle: get_or_create (create_rpc, wait_active with Dataproc state history, connect; or reuse with get_session, wait_endpoint, connect), stop (terminate_rpc, wait_terminating), and every GetSession (reason=rpc_guard|reuse|poll|stop). - every Spark Connect RPC at two layers of the channel stack: an outer span around the ProxiedChannel (what the client waits for, including the liveness guard, reported as guard_ns) and an inner span around the plain channel to the local proxy. Streaming calls go through an iterator proxy that keeps the grpc.Call surface, counts responses, ends the span at the final status timestamp, and handles abandoned streams via __del__ and weakref status callbacks. - the local proxy: start, each connection on the forwarding thread with byte/frame counters, and each websocket connect (ADC lookup, token refresh, handshake). - client-side conversions: plan building, Arrow to rows, Arrow to pandas. Spans nest through a contextvar and snapshot global attributes when they start, so a trace reads top-down from the user call to the wire. See docs/instrumentation.md for the API, the record schema and the span catalogue.
Recorder.inflight(prefix) lists spans that have started but not ended, and Recorder.wait_idle(timeout_s, prefix='grpc.') (also on Instrumentation) blocks until none match. Benchmarks use it to let the background ReleaseExecute RPCs pyspark issues after a query drain before the next measurement starts, instead of sleeping a fixed amount.
dborowitz
force-pushed
the
dborowitz-timing
branch
from
October 27, 2026 19:31
78b2072 to
ea38107
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.
Note
Part 1 of 2 in stack (followed by #211).
Summary
Adds an opt-in, dependency-free instrumentation package (
google.cloud.managed_spark_connect.instrumentation) that records timed spans, events, counters, and environment metadata across the Managed Spark Connect client stack.Instrumentation is off by default and adds zero overhead when disabled (all internal hooks route through a shared
NOOP_RECORDER, no gRPC interceptors are installed, and the PySpark client is left untouched). Users or benchmarks can opt in programmatically viaManagedSparkSession.builder.instrumentation(inst)or without code changes via theMANAGED_SPARK_CONNECT_INSTRUMENTATION_FILEandMANAGED_SPARK_CONNECT_INSTRUMENTATION_LOGenvironment variables.What is recorded
session.*,dataproc.get_session):session.get_or_create(session.create_rpc,session.wait_activewith server-sidestate_history,session.connect; orsession.reusewithdataproc.get_session,session.wait_endpoint,session.connect),session.stop(session.terminate_rpc,session.wait_terminating), and every control-planeGetSessionRPC tagged withreason=rpc_guard|reuse|poll|stopandnew_client.grpc.<Method>):layer="outer"wrapsProxiedChannel(measuring total client-perceived RPC latency, includingguard_nsspent in the session liveness guard).layer="inner"wraps the underlying channel to the local WebSocket proxy (measuring wire + server time).ExecutePlan,ReattachExecute) use a call-preserving iterator wrapper that tracks message/byte counts, recordsfirst_responseandstatusevents, freezesduration_nswhen the final gRPC status arrives on the channel thread, and safely emits spans even when PySpark drops an iterator atresult_completewithout draining it.proxy.*,auth.*,ws.handshake): Recordsproxy.start, per-connection forwarding spans (proxy.connectionwithbytes_in,bytes_out,frames_in,frames_out), and WebSocket connection establishment (proxy.ws_connect->auth.default,auth.refresh,ws.handshake).client.*): Tracksclient.plan_to_proto(plan_bytes),client.arrow_to_rows(rows), andclient.to_pandas->client.arrow_to_pandas(rows,columns).inflight(),wait_idle()): Allows callers and benchmarks to wait for asynchronous background gRPC calls (such as PySpark's post-queryReleaseExecuteRPCs) to finish before starting the next measurement.See
docs/instrumentation.mdfor the full API reference, JSONL record schema, and span catalogue.Commits in this PR
chore: ignore .venv/— Adds.venv/to.gitignore.fix: keep with_call()/future() on guarded ProxiedChannel methods— Replaces the plain function wrapper inProxiedChannel._wrap_methodwith_GuardedMultiCallable, preserving thewith_call()andfuture()methods required bygrpc.intercept_channel()and forwardingsubscribe()/unsubscribe()directly.fix: honor add_interceptor() on ManagedSparkChannelBuilder— Applies interceptors registered viaChannelBuilder.add_interceptor()around theProxiedChannelinManagedSparkChannelBuilder.toChannel().feat: opt-in instrumentation of sessions, RPCs, proxy and client— Implementsgoogle.cloud.managed_spark_connect.instrumentation(Instrumentation,Recorder,Span,JsonlFileSink,InMemorySink,LoggingSink,CallbackSink, gRPC interceptors, proxy/session/client hooks) and addsdocs/instrumentation.md.feat(instrumentation): track in-flight spans; add wait_idle()— AddsRecorder.inflight(prefix="")andInstrumentation.wait_idle(timeout_s, prefix="grpc.").Testing
tests/unit/test_instrumentation.py(44 tests),tests/unit/test_core.py, andtests/unit/test_proxy.pycovering span nesting, attribute snapshotting, streaming call lifecycle, abandoned stream finalization,with_call()/future()interceptor compatibility, andwait_idle().