diff --git a/.gitignore b/.gitignore index a7c34230..d1d12304 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,4 @@ __pycache__ build/ dist/ +.venv/ diff --git a/README.md b/README.md index 32803d27..97508091 100644 --- a/README.md +++ b/README.md @@ -116,6 +116,7 @@ The `ManagedSparkSession.builder` provides a fluent API to configure the session | `sessionTemplate(profile)` | Sets the Session Template to use. | | `subnetwork(subnet)` | Sets the subnetwork URI for the session. | | `ttl(duration)` | Sets the time-to-live (TTL) for the session using a `datetime.timedelta` object. | +| `instrumentation(inst)` | Records what the library does on your behalf; see [Instrumentation (opt-in)](#instrumentation-opt-in). | ### Reusing Named Sessions Across Notebooks @@ -149,6 +150,19 @@ To create or connect to a named session: 5. If a session with the specified ID is already in a `TERMINATED` or `FAILED` state (or is `TERMINATING`), the client emits a warning, deletes the old session, and creates a new session with the same ID. Deleting the old session removes its previous metadata and Spark UI history from Dataproc, though logs remain available in Cloud Logging. +### Instrumentation (opt-in) + +The library can record what it does on your behalf — Dataproc session +lifecycle steps, every Spark Connect gRPC call (with and without the local +proxy's overhead), the proxy's own connections, and the client-side plan and +result conversions — as timed *spans*. It is off by default; opt in via +`ManagedSparkSession.builder.instrumentation(...)` or the +`MANAGED_SPARK_CONNECT_INSTRUMENTATION_FILE` and +`MANAGED_SPARK_CONNECT_INSTRUMENTATION_LOG` environment variables. + +See [docs/instrumentation.md](https://github.com/GoogleCloudDataproc/managed-spark-connect-python/blob/main/docs/instrumentation.md) +for usage, the API, the record schema, and the span catalogue. + ### Interactive Extras For notebooks and other interactive use, include the `interactive` extra (along diff --git a/docs/instrumentation.md b/docs/instrumentation.md new file mode 100644 index 00000000..405a4fe3 --- /dev/null +++ b/docs/instrumentation.md @@ -0,0 +1,118 @@ +# Instrumentation (opt-in) + +The library can record what it does on your behalf — Dataproc session +lifecycle steps, every Spark Connect gRPC call (with and without the local +proxy's overhead), the proxy's own connections and the client-side plan and +result conversions — as timed *spans*. It is off by default and adds no work +to a session that has not opted in. + +Opt in from code: + +```python +from google.cloud.managed_spark_connect import ManagedSparkSession +from google.cloud.managed_spark_connect.instrumentation import ( + Instrumentation, JsonlFileSink, InMemorySink, +) + +inst = Instrumentation( + sinks=[JsonlFileSink("spark-connect-events.jsonl"), InMemorySink()], + attributes={"run_id": "nightly-42"}, # attached to every record +) +spark = ( + ManagedSparkSession.builder + .projectId(project).location(region) + .instrumentation(inst) + .getOrCreate() +) +spark.range(10).collect() + +spark.instrumentation is inst # True +inst.records # records of the InMemorySink(s) +inst.set_attribute("step", "load") # tag subsequent records +spark.stop() # flushes and closes the sinks +``` + +Or from the environment, without code changes (the builder wins when both +are used): + +| Variable | Effect | +|---|---| +| `MANAGED_SPARK_CONNECT_INSTRUMENTATION_FILE=/path/events.jsonl` | Append one JSON record per line to the file. | +| `MANAGED_SPARK_CONNECT_INSTRUMENTATION_LOG=1` | Log every record at `DEBUG` level through the `google.cloud.managed_spark_connect.instrumentation` logger. | + +## API + +```python +class Instrumentation: + def __init__(self, sinks=(), attributes=None, *, wrap_client=True) + @classmethod from_env() -> Instrumentation | None + recorder: Recorder; sinks: list[Sink] + set_attribute(key, value); remove_attribute(key) # global attrs on every subsequent record + records -> list[dict] # concatenated InMemorySink records + wrap_channel(channel, layer="outer") -> grpc.Channel + wait_idle(timeout_s, prefix="grpc.") -> bool # block until no such span is in flight (e.g. background ReleaseExecute) + flush(); close() # idempotent; file sinks reopen on the next write + +class Recorder: # thread-safe + span(name, **attrs) -> ContextManager[Span] # nests under the current span (contextvars) + start_span(name, parent=None, **attrs) -> Span # for threads that own their spans; call end() + event(name, **attrs); counter(name, delta=1, **attrs) + inflight(prefix="") -> list[Span]; wait_idle(timeout_s, prefix="grpc.") -> bool + trace_id: str; enabled: bool + +class Span: + span_id, parent_id, trace_id, name + set(**attrs); add_event(name, **attrs); record_exception(exc); end(end_mono_ns=None) + +# Sinks: write(record) / flush() / close() +JsonlFileSink(path, *, append=True) InMemorySink() LoggingSink(logger=None, level=DEBUG) CallbackSink(fn) + +# Lower level: GrpcInterceptor(recorder, layer), wrap_channel(channel, inst_or_recorder, layer), +# timed_guard(callback), install_client_wrappers(client, recorder) +``` + +`wrap_client=False` leaves the pyspark client untouched (no `client.*` spans). + +## Record schema + +Every record is a flat JSON object with a `type`. Global attributes are +merged into `attrs`. + +```json +{"type":"span","name":"grpc.ExecutePlan","span_id":"16 hex","parent_id":"16 hex | null", + "trace_id":"", + "start_wall_ns":0,"start_mono_ns":0,"duration_ns":0, + "status":"ok | error","error":{"type":"","message":"","grpc_code":"UNAVAILABLE"} | null, + "thread":"MainThread","attrs":{},"events":[{"name":"first_response","offset_ns":0}]} +{"type":"event","name":"...","wall_ns":0,"mono_ns":0,"thread":"...","attrs":{}} +{"type":"counter","name":"...","delta":1,"wall_ns":0,"attrs":{}} +{"type":"meta","wall_ns":0,"attrs":{"library_version":"","pyspark_version":"","grpcio_version":"", + "pyarrow_version":"","python":"","platform":"","pid":0}} +``` + +`start_wall_ns` is `time.time_ns()`; `start_mono_ns`, `duration_ns` and +event offsets come from `time.perf_counter_ns()`. One `meta` record is +written when an `Instrumentation` is created. Spans started on threads the +library owns (proxy connections, reattach pool threads) have +`parent_id: null` and the thread's name in `thread`. A span's `attrs` carry +the global attributes as they were when the span *started* (so a background +RPC stays attributed to the step that issued it); events, counters and the +meta record carry them as of when they are written. + +## Spans + +| Span | Attributes | +|---|---| +| `session.get_or_create` → `session.create_rpc`, `session.wait_active`, `session.connect` (create) or `session.reuse` → `dataproc.get_session`, `session.wait_endpoint`, `session.connect` (reuse) | `mode=create\|reuse`, `session_id`, `uuid`, `runtime_version`, `template`; `session.wait_active`/`session.wait_endpoint` add `poll_interval_s`, `create_time`, `state_time`, `state`, `state_history` (JSON string); `session.reuse` adds `found` | +| `session.stop` → `session.terminate_rpc`, `session.wait_terminating` | `terminate=true\|false`, `session_id`; `session.wait_terminating` adds `state` | +| `dataproc.get_session` | `reason=rpc_guard\|reuse\|poll\|stop`, `new_client` (whether a fresh `SessionControllerClient` was built for the call) | +| `grpc.` (e.g. `grpc.ExecutePlan`), one `layer=outer` span (what the client waits for, including the liveness guard) wrapping one `layer=inner` span (pure wire time to the proxy) | `method`, `kind=unary_unary\|unary_stream\|stream_unary\|stream_stream`, `layer`, `session_id`, `operation_id`, `request_bytes`, `request_count` (client streaming), `response_count`, `response_bytes`, `grpc_code`, `guard_ns` (outer only); events `first_response`, `status`. Byte counts are protobuf `ByteSize()` of the decoded messages, not wire bytes (the proxy's websocket framing roughly doubles those); a streaming span counts the messages delivered to the consumer, freezes its `duration_ns` when the final gRPC status arrives, and emits the record when the consumer is done with the stream (iteration finished, error, `cancel()`, or the iterator being dropped — pyspark drops an `ExecutePlan` stream at `result_complete` without draining it; that last case relies on CPython's reference counting to be prompt) | +| `proxy.start` | `port` | +| `proxy.connection` (accept to close, on the forwarding thread) → `proxy.ws_connect` → `auth.default`, `auth.refresh`, `ws.handshake` | `conn`, `bytes_in`, `bytes_out`, `frames_in`, `frames_out` (wire side, i.e. the raw TCP bytes forwarded); `proxy.ws_connect` has `conn`, `host` | +| `client.plan_to_proto` | `plan_bytes` | +| `client.arrow_to_rows` | `rows` (recorded to the most recently instrumented session when several share a process) | +| `client.to_pandas` → `client.arrow_to_pandas` | `rows`, `columns` on the conversion span | + +The outer `grpc.*` span is also the parent of the `dataproc.get_session` +span the liveness guard issues, so for an RPC `outer − inner − guard_ns` is +the client's own overhead and `inner` is time on the wire plus server time. diff --git a/google/cloud/managed_spark_connect/client/core.py b/google/cloud/managed_spark_connect/client/core.py index 843e741c..159de2a9 100644 --- a/google/cloud/managed_spark_connect/client/core.py +++ b/google/cloud/managed_spark_connect/client/core.py @@ -18,6 +18,7 @@ from pyspark.sql.connect.client import DefaultChannelBuilder from . import proxy +from .. import instrumentation as instrumentation_lib logger = logging.getLogger(__name__) @@ -40,8 +41,9 @@ class ManagedSparkChannelBuilder(DefaultChannelBuilder): True """ - def __init__(self, url, is_active_callback=None): + def __init__(self, url, is_active_callback=None, instrumentation=None): self._is_active_callback = is_active_callback + self._instrumentation = instrumentation super().__init__(url) def toChannel(self) -> grpc.Channel: @@ -59,7 +61,17 @@ def toChannel(self) -> grpc.Channel: return self._proxied_channel() def _proxied_channel(self) -> grpc.Channel: - return ProxiedChannel(self.host, self._is_active_callback) + channel = ProxiedChannel( + self.host, self._is_active_callback, self._instrumentation + ) + if self._interceptors: + logger.debug(f"Applying interceptors ({self._interceptors})") + channel = grpc.intercept_channel(channel, *self._interceptors) + # Outermost, so the span covers the user's interceptors and the + # liveness guard as well as the RPC itself. + return instrumentation_lib.wrap_channel( + channel, self._instrumentation, instrumentation_lib.OUTER + ) def _direct_channel(self) -> grpc.Channel: destination = f"{self.host}:{self.port}" @@ -81,16 +93,60 @@ def _direct_channel(self) -> grpc.Channel: ) +class _GuardedMultiCallable: + """A gRPC multi-callable that runs a guard before every invocation. + + ``grpc.Channel.unary_unary()`` and friends return multi-callables that + are invoked either directly (``stub.Method(request)``), through + ``with_call()`` (as ``grpc.intercept_channel`` does for unary-unary and + stream-unary RPCs) or through ``future()``. All three entry points are + forwarded here so a guarded :class:`ProxiedChannel` stays a drop-in + replacement for the channel it wraps. + """ + + def __init__(self, wrapped_method, guard): + self._wrapped_method = wrapped_method + self._guard = guard + + def __call__(self, *args, **kwargs): + self._guard() + return self._wrapped_method(*args, **kwargs) + + def with_call(self, *args, **kwargs): + self._guard() + return self._wrapped_method.with_call(*args, **kwargs) + + def future(self, *args, **kwargs): + self._guard() + return self._wrapped_method.future(*args, **kwargs) + + class ProxiedChannel(grpc.Channel): + """A gRPC channel to a Spark Connect server behind the local proxy. + + Every RPC first runs ``is_active_callback`` (when given) and is rejected + with a ``RuntimeError`` if the session is gone. With ``instrumentation`` + the callback is timed and the inner, pure-wire channel to the proxy is + traced as well (see the ``instrumentation`` package). + """ - def __init__(self, target_host, is_active_callback): + def __init__(self, target_host, is_active_callback, instrumentation=None): + recorder = instrumentation_lib.get_recorder(instrumentation) + if recorder.enabled: + is_active_callback = instrumentation_lib.timed_guard( + is_active_callback + ) self._is_active_callback = is_active_callback - self._proxy = proxy.ManagedSparkSessionProxy(0, target_host) + self._proxy = proxy.ManagedSparkSessionProxy( + 0, target_host, recorder=recorder + ) self._proxy.start() self._proxied_connect_url = f"sc://localhost:{self._proxy.port}" - self._wrapped = DefaultChannelBuilder( - self._proxied_connect_url - ).toChannel() + self._wrapped = instrumentation_lib.wrap_channel( + DefaultChannelBuilder(self._proxied_connect_url).toChannel(), + recorder, + instrumentation_lib.INNER, + ) def __enter__(self): return self @@ -105,22 +161,20 @@ def close(self): self._proxy.stop() return ret + def _check_active(self): + if ( + self._is_active_callback is not None + and not self._is_active_callback() + ): + logger.warning(f"Session is no longer active") + raise RuntimeError( + "Session not active. Please create a new session" + ) + def _wrap_method(self, wrapped_method): if self._is_active_callback is None: return wrapped_method - - def checked_method(*margs, **mkwargs): - if ( - self._is_active_callback is not None - and not self._is_active_callback() - ): - logger.warning(f"Session is no longer active") - raise RuntimeError( - "Session not active. Please create a new session" - ) - return wrapped_method(*margs, **mkwargs) - - return checked_method + return _GuardedMultiCallable(wrapped_method, self._check_active) def stream_stream(self, *args, **kwargs): return self._wrap_method(self._wrapped.stream_stream(*args, **kwargs)) @@ -129,7 +183,7 @@ def stream_unary(self, *args, **kwargs): return self._wrap_method(self._wrapped.stream_unary(*args, **kwargs)) def subscribe(self, *args, **kwargs): - return self._wrap_method(self._wrapped.subscribe(*args, **kwargs)) + return self._wrapped.subscribe(*args, **kwargs) def unary_stream(self, *args, **kwargs): return self._wrap_method(self._wrapped.unary_stream(*args, **kwargs)) @@ -138,4 +192,4 @@ def unary_unary(self, *args, **kwargs): return self._wrap_method(self._wrapped.unary_unary(*args, **kwargs)) def unsubscribe(self, *args, **kwargs): - return self._wrap_method(self._wrapped.unsubscribe(*args, **kwargs)) + return self._wrapped.unsubscribe(*args, **kwargs) diff --git a/google/cloud/managed_spark_connect/client/proxy.py b/google/cloud/managed_spark_connect/client/proxy.py index cf680439..1d26852c 100755 --- a/google/cloud/managed_spark_connect/client/proxy.py +++ b/google/cloud/managed_spark_connect/client/proxy.py @@ -18,12 +18,15 @@ import logging import socket import threading +import time import websockets.sync.client as websocketclient from google import auth as googleauth from google.auth.transport import requests as googleauthrequests +from google.cloud.managed_spark_connect.instrumentation import NOOP_RECORDER + parser = argparse.ArgumentParser() parser.add_argument("port") parser.add_argument("target_host") @@ -59,7 +62,7 @@ def close(self): return self._conn.close() -def connect_tcp_bridge(hostname): +def connect_tcp_bridge(hostname, recorder=NOOP_RECORDER, conn_number=None): """Create a socket-like connection to the given hostname using websocket. The backend server connected to over the websocket connection must be @@ -67,24 +70,44 @@ def connect_tcp_bridge(hostname): Args: hostname: The hostname of the server running the TCP-bridge backend. + recorder: Instrumentation recorder for the ``proxy.ws_connect`` span. + conn_number: The proxy connection this websocket serves, if any. Returns: A socket-like object with `recv` and `send` methods. """ path = "tcp-over-websocket-bridge/35218cb7-1201-4940-89e8-48d8f03fed96" - creds, _ = googleauth.default( - scopes=["https://www.googleapis.com/auth/cloud-platform"] - ) - creds.refresh(googleauthrequests.Request()) + with recorder.span("proxy.ws_connect", conn=conn_number, host=hostname): + with recorder.span("auth.default"): + creds, _ = googleauth.default( + scopes=["https://www.googleapis.com/auth/cloud-platform"] + ) + with recorder.span("auth.refresh"): + creds.refresh(googleauthrequests.Request()) + + with recorder.span("ws.handshake"): + return websocketclient.connect( + f"wss://{hostname}/{path}", + additional_headers={"Authorization": f"Bearer {creds.token}"}, + open_timeout=30, + ) + + +class _DirectionStats: + """Byte and frame counters for one forwarding direction. + + Each direction is driven by exactly one thread, so plain attributes are + enough; the totals are read only after both threads have finished. + """ - return websocketclient.connect( - f"wss://{hostname}/{path}", - additional_headers={"Authorization": f"Bearer {creds.token}"}, - open_timeout=30, - ) + __slots__ = ("bytes", "frames") + + def __init__(self): + self.bytes = 0 + self.frames = 0 -def forward_bytes(name, from_sock, to_sock): +def forward_bytes(name, from_sock, to_sock, stats=None): """Continuously stream bytes from the `from_sock` to the `to_sock`. This method terminates when either the `from_sock` is closed (causing @@ -97,6 +120,7 @@ def forward_bytes(name, from_sock, to_sock): name: forwarding thread name from_sock: A socket-like object to stream bytes from. to_sock: A socket-like object to stream bytes to. + stats: Optional :class:`_DirectionStats` to count what was forwarded. """ while True: try: @@ -104,6 +128,9 @@ def forward_bytes(name, from_sock, to_sock): if not bs: to_sock.close() return + if stats is not None: + stats.bytes += len(bs) + stats.frames += 1 attempt = 0 while bs and (attempt < 10): attempt += 1 @@ -125,19 +152,24 @@ def forward_bytes(name, from_sock, to_sock): return -def connect_sockets(conn_number, from_sock, to_sock): +def connect_sockets(conn_number, from_sock, to_sock, span=None): """Create a connection between the two given ports. This method continuously streams bytes in both directions between the given `from_sock` and `to_sock` socket-like objects. The caller is responsible for creating and closing the supplied sockets. + When ``span`` is given, ``bytes_out``/``frames_out`` (``from_sock`` to + ``to_sock``) and ``bytes_in``/``frames_in`` (the other way) are set on + it once both directions have finished. """ + out_stats = _DirectionStats() if span is not None else None + in_stats = _DirectionStats() if span is not None else None forward_name = f"{conn_number}-forward" t1 = threading.Thread( name=forward_name, target=forward_bytes, - args=[forward_name, from_sock, to_sock], + args=[forward_name, from_sock, to_sock, out_stats], daemon=True, ) t1.start() @@ -145,15 +177,29 @@ def connect_sockets(conn_number, from_sock, to_sock): t2 = threading.Thread( name=backward_name, target=forward_bytes, - args=[backward_name, to_sock, from_sock], + args=[backward_name, to_sock, from_sock, in_stats], daemon=True, ) t2.start() t1.join() t2.join() - - -def forward_connection(conn_number, conn, addr, target_host): + if span is not None: + span.set( + bytes_in=in_stats.bytes, + bytes_out=out_stats.bytes, + frames_in=in_stats.frames, + frames_out=out_stats.frames, + ) + + +def forward_connection( + conn_number, + conn, + addr, + target_host, + recorder=NOOP_RECORDER, + accepted_ns=None, +): """Create a connection to the target and forward `conn` to it. This method creates a socket-like object holding a connection to the given @@ -165,16 +211,42 @@ def forward_connection(conn_number, conn, addr, target_host): This method should be run inside a daemon thread so that it will not block program termination. + + Args: + conn_number: Sequence number of this connection within the proxy. + conn: The accepted client socket. + addr: The client address (for logging). + target_host: The backend host to connect to. + recorder: Instrumentation recorder. A ``proxy.connection`` span + covers this call (its children: ``proxy.ws_connect`` and the + auth/handshake steps) and carries the traffic counters. + accepted_ns: ``(perf_counter_ns, time_ns)`` of the accept, so the + span starts when the client connected rather than when this + thread got scheduled. """ - with conn: - with connect_tcp_bridge(target_host) as websocket_conn: - backend_socket = bridged_socket(websocket_conn) - # Set a timeout on how long we will allow send/recv calls to block - # - # The code that reads and writes to this connection will retry - # on timeouts, so this is a safe change. - conn.settimeout(10) - connect_sockets(conn_number, conn, backend_socket) + if accepted_ns is None: + span = recorder.start_span("proxy.connection", conn=conn_number) + else: + span = recorder.start_span_at( + "proxy.connection", *accepted_ns, conn=conn_number + ) + with span: + with conn: + with connect_tcp_bridge( + target_host, recorder, conn_number + ) as websocket_conn: + backend_socket = bridged_socket(websocket_conn) + # Set a timeout on how long we will allow send/recv calls to block + # + # The code that reads and writes to this connection will retry + # on timeouts, so this is a safe change. + conn.settimeout(10) + connect_sockets( + conn_number, + conn, + backend_socket, + span if recorder.enabled else None, + ) class ManagedSparkSessionProxy(object): @@ -188,9 +260,10 @@ class ManagedSparkSessionProxy(object): Default Credentials. """ - def __init__(self, port, target_host): + def __init__(self, port, target_host, recorder=None): self._port = port self._target_host = target_host + self._recorder = recorder if recorder is not None else NOOP_RECORDER self._started = False self._killed = False self._conn_number = 0 @@ -209,10 +282,12 @@ def start(self, daemon=True): if self._started: raise Exception("Managed Spark session proxy already started") self._started = True - s = threading.Semaphore(value=0) - t = threading.Thread(target=self._run, args=[s], daemon=daemon) - t.start() - s.acquire() + with self._recorder.span("proxy.start") as span: + s = threading.Semaphore(value=0) + t = threading.Thread(target=self._run, args=[s], daemon=daemon) + t.start() + s.acquire() + span.set(port=self._port) def _run(self, s): with socket.create_server(("127.0.0.1", self._port)) as frontend_socket: @@ -221,11 +296,21 @@ def _run(self, s): s.release() while not self._killed: conn, addr = frontend_socket.accept() + accepted_ns = None + if self._recorder.enabled: + accepted_ns = (time.perf_counter_ns(), time.time_ns()) logger.debug(f"Accepted a connection from {addr}...") self._conn_number += 1 threading.Thread( target=forward_connection, - args=[self._conn_number, conn, addr, self._target_host], + args=[ + self._conn_number, + conn, + addr, + self._target_host, + self._recorder, + accepted_ns, + ], daemon=True, ).start() diff --git a/google/cloud/managed_spark_connect/instrumentation/__init__.py b/google/cloud/managed_spark_connect/instrumentation/__init__.py new file mode 100644 index 00000000..8a670108 --- /dev/null +++ b/google/cloud/managed_spark_connect/instrumentation/__init__.py @@ -0,0 +1,100 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Opt-in instrumentation for Managed Spark Connect sessions. + +Off by default. Turn it on by passing an :class:`Instrumentation` to the +session builder, or through the environment:: + + from google.cloud.managed_spark_connect import ManagedSparkSession + from google.cloud.managed_spark_connect.instrumentation import ( + Instrumentation, JsonlFileSink, InMemorySink, + ) + + inst = Instrumentation( + sinks=[JsonlFileSink("events.jsonl"), InMemorySink()], + attributes={"run_id": "2024-01-01-a"}, + ) + spark = ( + ManagedSparkSession.builder.projectId(p).location(r) + .instrumentation(inst) + .getOrCreate() + ) + spark.range(10).collect() + spark.instrumentation.records # the InMemorySink's records + + # or, without code changes: + # MANAGED_SPARK_CONNECT_INSTRUMENTATION_FILE=/tmp/events.jsonl + # MANAGED_SPARK_CONNECT_INSTRUMENTATION_LOG=1 + +Records are dicts with a ``type`` of ``span``, ``event``, ``counter`` or +``meta``. A span record looks like:: + + {"type": "span", "name": "grpc.ExecutePlan", "span_id": "16 hex", + "parent_id": "16 hex or null", "trace_id": "", + "start_wall_ns": ..., "start_mono_ns": ..., "duration_ns": ..., + "status": "ok|error", + "error": {"type": "", "message": "", "grpc_code": "UNAVAILABLE"} | null, + "thread": "MainThread", "attrs": {...}, + "events": [{"name": "first_response", "offset_ns": ...}]} + +``start_wall_ns`` is ``time.time_ns()``; ``start_mono_ns`` and every +duration/offset come from ``time.perf_counter_ns()``. Events and counters +carry ``wall_ns`` (and events ``mono_ns``) instead of a start/duration. The +single ``meta`` record per :class:`Instrumentation` lists library, pyspark, +grpcio, pyarrow, Python and platform versions. Global attributes set on the +:class:`Instrumentation` are merged into every record's ``attrs``. + +See ``docs/instrumentation.md`` for the span catalogue. +""" + +from ._client import install_client_wrappers, uninstall_client_wrappers +from ._core import ( + ENV_FILE, + ENV_LOG, + NOOP_RECORDER, + Instrumentation, + Recorder, + Span, + get_recorder, +) +from ._grpc import INNER, OUTER, GrpcInterceptor, timed_guard, wrap_channel +from ._sinks import ( + CallbackSink, + InMemorySink, + JsonlFileSink, + LoggingSink, + Sink, +) + +__all__ = [ + "CallbackSink", + "ENV_FILE", + "ENV_LOG", + "GrpcInterceptor", + "INNER", + "InMemorySink", + "Instrumentation", + "JsonlFileSink", + "LoggingSink", + "NOOP_RECORDER", + "OUTER", + "Recorder", + "Sink", + "Span", + "get_recorder", + "install_client_wrappers", + "timed_guard", + "uninstall_client_wrappers", + "wrap_channel", +] diff --git a/google/cloud/managed_spark_connect/instrumentation/_client.py b/google/cloud/managed_spark_connect/instrumentation/_client.py new file mode 100644 index 00000000..bb4370d3 --- /dev/null +++ b/google/cloud/managed_spark_connect/instrumentation/_client.py @@ -0,0 +1,162 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Spans for the client-side work around an RPC: plan building and result +conversion. + +Nothing here is touched unless an instrumented session is created. Then: + +* ``client.to_pandas`` and ``client.arrow_to_pandas`` come from wrappers + installed on that one ``SparkConnectClient`` instance (``to_pandas`` and + ``_execute_and_fetch``): the conversion span starts when the Arrow table + has been fetched and ends when ``to_pandas`` returns. +* ``client.plan_to_proto`` wraps ``LogicalPlan.to_proto`` at class level, + because pyspark builds the plan on the DataFrame before handing it to the + client. The wrapper dispatches on the client it is given: only clients + that carry a recorder get a span, others pay one attribute lookup. +* ``client.arrow_to_rows`` wraps ``ArrowTableToRowsConversion.convert`` at + class level. That static method knows no client, so it records to the + most recently instrumented client's recorder (a process has at most one + managed session at a time). + +The class-level hooks are installed once per process, on first use, and +are inert while no instrumented client is registered. +""" + +import contextvars +import threading +from typing import Any, List, Optional + +from ._core import NOOP_RECORDER, Recorder + +_RECORDER_ATTR = "_managed_spark_connect_recorder" + +_lock = threading.Lock() +_registered: List[Recorder] = [] +_hooks_installed = False + +# Set while an instrumented ``to_pandas`` runs: [to_pandas span, conversion +# span or None]. Lets the ``_execute_and_fetch`` wrapper start the +# conversion span exactly when the fetch finishes. +_TO_PANDAS: contextvars.ContextVar[Optional[list]] = contextvars.ContextVar( + "managed_spark_connect_to_pandas", default=None +) + + +def _active_recorder() -> Recorder: + with _lock: + return _registered[-1] if _registered else NOOP_RECORDER + + +def _recorder_of(client: Any) -> Recorder: + return getattr(client, _RECORDER_ATTR, None) or NOOP_RECORDER + + +def _install_class_hooks() -> None: + global _hooks_installed + with _lock: + if _hooks_installed: + return + _hooks_installed = True + + from pyspark.sql.connect.plan import LogicalPlan + from pyspark.sql.conversion import ArrowTableToRowsConversion + + original_to_proto = LogicalPlan.to_proto + + def to_proto(self, session, *args, **kwargs): + recorder = _recorder_of(session) + if not recorder.enabled: + return original_to_proto(self, session, *args, **kwargs) + with recorder.span("client.plan_to_proto") as span: + plan = original_to_proto(self, session, *args, **kwargs) + span.set(plan_bytes=plan.ByteSize()) + return plan + + to_proto.__wrapped__ = original_to_proto # type: ignore[attr-defined] + LogicalPlan.to_proto = to_proto + + original_convert = ArrowTableToRowsConversion.convert + + def convert(table, schema, *args, **kwargs): + recorder = _active_recorder() + if not recorder.enabled: + return original_convert(table, schema, *args, **kwargs) + with recorder.span("client.arrow_to_rows", rows=table.num_rows): + return original_convert(table, schema, *args, **kwargs) + + convert.__wrapped__ = original_convert # type: ignore[attr-defined] + ArrowTableToRowsConversion.convert = staticmethod(convert) + + +def install_client_wrappers(client: Any, recorder: Recorder) -> None: + """Instrument one ``SparkConnectClient`` instance with ``recorder``.""" + if not recorder.enabled or getattr(client, _RECORDER_ATTR, None): + return + _install_class_hooks() + setattr(client, _RECORDER_ATTR, recorder) + with _lock: + _registered.append(recorder) + + original_to_pandas = client.to_pandas + original_execute_and_fetch = client._execute_and_fetch + + def to_pandas(*args, **kwargs): + with recorder.span("client.to_pandas") as span: + state = [span, None] + token = _TO_PANDAS.set(state) + try: + return original_to_pandas(*args, **kwargs) + except BaseException as e: + if state[1] is not None: + state[1].record_exception(e) + raise + finally: + _TO_PANDAS.reset(token) + if state[1] is not None: + state[1].end() + + def _execute_and_fetch(*args, **kwargs): + result = original_execute_and_fetch(*args, **kwargs) + state = _TO_PANDAS.get() + if state is not None and state[1] is None: + table = result[0] if isinstance(result, tuple) else None + state[1] = recorder.start_span( + "client.arrow_to_pandas", + parent=state[0], + rows=getattr(table, "num_rows", None), + columns=getattr(table, "num_columns", None), + ) + return result + + to_pandas.__wrapped__ = original_to_pandas # type: ignore[attr-defined] + _execute_and_fetch.__wrapped__ = original_execute_and_fetch # type: ignore[attr-defined] + client.to_pandas = to_pandas + client._execute_and_fetch = _execute_and_fetch + + +def uninstall_client_wrappers(client: Any) -> None: + """Undo :func:`install_client_wrappers` for ``client`` (if applied).""" + recorder = getattr(client, _RECORDER_ATTR, None) + if recorder is None: + return + for name in ("to_pandas", "_execute_and_fetch"): + wrapper = client.__dict__.get(name) + if wrapper is not None and hasattr(wrapper, "__wrapped__"): + del client.__dict__[name] + delattr(client, _RECORDER_ATTR) + with _lock: + try: + _registered.remove(recorder) + except ValueError: + pass diff --git a/google/cloud/managed_spark_connect/instrumentation/_core.py b/google/cloud/managed_spark_connect/instrumentation/_core.py new file mode 100644 index 00000000..666fcac9 --- /dev/null +++ b/google/cloud/managed_spark_connect/instrumentation/_core.py @@ -0,0 +1,615 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Spans, the recorder that creates them and the ``Instrumentation`` handle. + +Everything here is deliberately small and dependency free: a span is a +plain object that turns into a ``dict`` when it ends, the recorder tracks +the current span with a :class:`contextvars.ContextVar` so nesting works +across ``with`` blocks (and across threads whenever the context is copied), +and the no-op recorder used while instrumentation is off never allocates. +""" + +import contextvars +import importlib.metadata +import logging +import os +import platform +import random +import sys +import threading +import time +from typing import ( + Any, + Callable, + ContextManager, + Dict, + List, + Mapping, + Optional, + Sequence, +) + +from ._sinks import ( + CallbackSink, + InMemorySink, + JsonlFileSink, + LoggingSink, + Sink, +) + +logger = logging.getLogger(__name__) + +ENV_FILE = "MANAGED_SPARK_CONNECT_INSTRUMENTATION_FILE" +ENV_LOG = "MANAGED_SPARK_CONNECT_INSTRUMENTATION_LOG" + +_CURRENT_SPAN: contextvars.ContextVar[Optional["Span"]] = ( + contextvars.ContextVar("managed_spark_connect_current_span", default=None) +) + +_TRUE_VALUES = {"1", "true", "yes", "on"} + + +def _new_id() -> str: + return f"{random.getrandbits(64):016x}" + + +def _grpc_code_name(exc: BaseException) -> Optional[str]: + """Best-effort gRPC status code name for ``exc``. + + Handles ``grpc.RpcError`` (which carries ``code()``) and + ``google.api_core`` exceptions (which carry ``grpc_status_code``). + """ + code = None + code_fn = getattr(exc, "code", None) + if callable(code_fn): + try: + code = code_fn() + except Exception: # pragma: no cover - defensive + code = None + if code is None: + code = getattr(exc, "grpc_status_code", None) + if code is None: + return None + return getattr(code, "name", str(code)) + + +def error_info(exc: BaseException) -> Dict[str, Any]: + """The ``error`` object recorded for an exception.""" + message = None + details = getattr(exc, "details", None) + if callable(details): + # grpc.RpcError: the status message, not the multi-line repr. + try: + message = details() + except Exception: # pragma: no cover - defensive + message = None + if not isinstance(message, str): + message = str(exc) + return { + "type": type(exc).__name__, + "message": message, + "grpc_code": _grpc_code_name(exc), + } + + +class Span: + """A timed, named unit of work. + + Use it as a context manager (``with recorder.span(...) as span:``) to + have it become the parent of spans opened inside the block, or call + :meth:`end` yourself for spans that outlive a single frame (proxy + connections, streaming RPCs). + """ + + __slots__ = ( + "name", + "span_id", + "parent_id", + "trace_id", + "attrs", + "events", + "start_wall_ns", + "start_mono_ns", + "thread", + "status", + "error", + "_recorder", + "_ended", + "_lock", + "_token", + ) + + def __init__( + self, + recorder: "Recorder", + name: str, + parent_id: Optional[str], + attrs: Dict[str, Any], + *, + start_mono_ns: Optional[int] = None, + start_wall_ns: Optional[int] = None, + ): + self._recorder = recorder + self.name = name + self.span_id = _new_id() + self.parent_id = parent_id + self.trace_id = recorder.trace_id + # Global attributes are snapshotted now, not when the span ends, so a + # span is attributed to the trial/step that started it even if it + # outlives it (background ReleaseExecute RPCs, proxy connections). + self.attrs = recorder._attributes() + self.attrs.update(attrs) + self.events: List[Dict[str, Any]] = [] + self.start_mono_ns = ( + time.perf_counter_ns() if start_mono_ns is None else start_mono_ns + ) + self.start_wall_ns = ( + time.time_ns() if start_wall_ns is None else start_wall_ns + ) + self.thread = threading.current_thread().name + self.status = "ok" + self.error: Optional[Dict[str, Any]] = None + self._ended = False + self._lock = threading.Lock() + self._token = None + recorder._track(self) + + def set(self, **attrs: Any) -> None: + """Add or overwrite attributes.""" + with self._lock: + if not self._ended: + self.attrs.update(attrs) + + def add_event(self, name: str, **attrs: Any) -> None: + """Record a point in time inside the span (``offset_ns`` is added).""" + event = { + "name": name, + "offset_ns": time.perf_counter_ns() - self.start_mono_ns, + } + if attrs: + event.update(attrs) + with self._lock: + if not self._ended: + self.events.append(event) + + def record_exception(self, exc: BaseException) -> None: + """Mark the span failed; the error keeps type, message, grpc code.""" + info = error_info(exc) + with self._lock: + if self._ended: + return + self.status = "error" + self.error = info + + @property + def ended(self) -> bool: + return self._ended + + def end(self, end_mono_ns: Optional[int] = None) -> None: + """Finish the span and hand its record to the sinks. Idempotent. + + ``end_mono_ns`` (a ``perf_counter_ns`` value) freezes the end at an + earlier instant than now, for spans whose end was observed before + they could be emitted (a streaming RPC whose status arrived while + the consumer still held the iterator). + """ + if end_mono_ns is None: + end_mono_ns = time.perf_counter_ns() + with self._lock: + if self._ended: + return + self._ended = True + self._recorder._untrack(self) + self._recorder._emit( + { + "type": "span", + "name": self.name, + "span_id": self.span_id, + "parent_id": self.parent_id, + "trace_id": self.trace_id or self._recorder.trace_id, + "start_wall_ns": self.start_wall_ns, + "start_mono_ns": self.start_mono_ns, + "duration_ns": max(0, end_mono_ns - self.start_mono_ns), + "status": self.status, + "error": self.error, + "thread": self.thread, + "attrs": self.attrs, + "events": self.events, + } + ) + + def __enter__(self) -> "Span": + self._token = _CURRENT_SPAN.set(self) + return self + + def __exit__(self, exc_type, exc, tb) -> bool: + if exc is not None: + self.record_exception(exc) + token, self._token = self._token, None + if token is not None: + try: + _CURRENT_SPAN.reset(token) + except ValueError: + # Exited in a different context than it was entered in. + _CURRENT_SPAN.set(None) + self.end() + return False + + def __repr__(self) -> str: + return ( + f"Span({self.name!r}, id={self.span_id}, parent={self.parent_id})" + ) + + +class _NoopSpan: + """The span handed out while instrumentation is disabled. + + One shared instance; every method is a no-op, so ``with + recorder.span(...)`` costs a method call and nothing else. + """ + + __slots__ = () + + name = "" + span_id = "" + parent_id = None + trace_id = "" + attrs: Dict[str, Any] = {} + events: List[Dict[str, Any]] = [] + ended = True + + def set(self, **attrs: Any) -> None: + pass + + def add_event(self, name: str, **attrs: Any) -> None: + pass + + def record_exception(self, exc: BaseException) -> None: + pass + + def end(self, end_mono_ns: Optional[int] = None) -> None: + pass + + def __enter__(self) -> "_NoopSpan": + return self + + def __exit__(self, exc_type, exc, tb) -> bool: + return False + + +_NOOP_SPAN = _NoopSpan() + + +class Recorder: + """Creates spans, events and counters. Thread-safe. + + ``span()`` nests under whatever span is current in the calling context; + ``start_span()`` does not look at the context (pass ``parent`` + explicitly), which is what threads that own their own spans want. + """ + + enabled = True + + def __init__( + self, + emit: Callable[[Dict[str, Any]], None], + attributes: Callable[[], Dict[str, Any]] = dict, + ): + self._emit = emit + self._attributes = attributes + self.trace_id = "" + self._inflight: Dict[str, "Span"] = {} + self._idle = threading.Condition() + + def current_span(self) -> Optional[Span]: + """The span that ``span()`` would nest under, if any.""" + return _CURRENT_SPAN.get() + + def inflight(self, prefix: str = "") -> List[Span]: + """Spans started but not yet ended whose name starts with ``prefix``.""" + with self._idle: + return [ + s for s in self._inflight.values() if s.name.startswith(prefix) + ] + + def wait_idle(self, timeout_s: float, prefix: str = "grpc.") -> bool: + """Block until no span named ``prefix*`` is in flight. + + Returns ``False`` if ``timeout_s`` elapsed first. The default prefix + waits for RPCs, including the background ``ReleaseExecute`` calls + pyspark issues on its thread pool after a query finishes. + """ + deadline = time.monotonic() + timeout_s + with self._idle: + while any( + s.name.startswith(prefix) for s in self._inflight.values() + ): + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + self._idle.wait(remaining) + return True + + def _track(self, span: "Span") -> None: + with self._idle: + self._inflight[span.span_id] = span + + def _untrack(self, span: "Span") -> None: + with self._idle: + if self._inflight.pop(span.span_id, None) is not None: + self._idle.notify_all() + + def span(self, name: str, **attrs: Any) -> ContextManager[Span]: + parent = _CURRENT_SPAN.get() + return Span(self, name, parent.span_id if parent else None, attrs) + + def start_span( + self, name: str, parent: Optional[Span] = None, **attrs: Any + ) -> Span: + return Span(self, name, parent.span_id if parent else None, attrs) + + def start_span_at( + self, + name: str, + start_mono_ns: int, + start_wall_ns: int, + parent: Optional[Span] = None, + **attrs: Any, + ) -> Span: + """Like :meth:`start_span`, for work that began before the span could + be created (``start_*_ns`` as ``perf_counter_ns``/``time_ns``).""" + return Span( + self, + name, + parent.span_id if parent else None, + attrs, + start_mono_ns=start_mono_ns, + start_wall_ns=start_wall_ns, + ) + + def event(self, name: str, **attrs: Any) -> None: + self._emit( + { + "type": "event", + "name": name, + "wall_ns": time.time_ns(), + "mono_ns": time.perf_counter_ns(), + "thread": threading.current_thread().name, + "attrs": attrs, + } + ) + + def counter(self, name: str, delta: float = 1, **attrs: Any) -> None: + self._emit( + { + "type": "counter", + "name": name, + "delta": delta, + "wall_ns": time.time_ns(), + "attrs": attrs, + } + ) + + +class _NoopRecorder: + """Drop-in for :class:`Recorder` that records nothing.""" + + enabled = False + + @property + def trace_id(self) -> str: + return "" + + @trace_id.setter + def trace_id(self, value: str) -> None: + pass + + def current_span(self) -> None: + return None + + def inflight(self, prefix: str = "") -> List[Span]: + return [] + + def wait_idle(self, timeout_s: float, prefix: str = "grpc.") -> bool: + return True + + def span(self, name: str, **attrs: Any) -> ContextManager[_NoopSpan]: + return _NOOP_SPAN + + def start_span(self, name: str, parent=None, **attrs: Any) -> _NoopSpan: + return _NOOP_SPAN + + def start_span_at( + self, name, start_mono_ns, start_wall_ns, parent=None, **attrs + ): + return _NOOP_SPAN + + def event(self, name: str, **attrs: Any) -> None: + pass + + def counter(self, name: str, delta: float = 1, **attrs: Any) -> None: + pass + + +NOOP_RECORDER = _NoopRecorder() +# The name library code reaches for when it has no Instrumentation. +_NOOP_RECORDER = NOOP_RECORDER + + +def _version_of(distribution: str) -> Optional[str]: + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + return None + + +def _pyspark_version() -> Optional[str]: + try: + from pyspark.version import __version__ + + return __version__ + except Exception: + return None + + +def _meta_attrs() -> Dict[str, Any]: + return { + "library_version": _version_of("google-cloud-spark-connect"), + "pyspark_version": _pyspark_version(), + "grpcio_version": _version_of("grpcio"), + "pyarrow_version": _version_of("pyarrow"), + "python": platform.python_version(), + "platform": platform.platform(), + "pid": os.getpid(), + } + + +class Instrumentation: + """Opt-in tracing of what the library does on the user's behalf. + + Pass an instance to ``ManagedSparkSession.builder.instrumentation()`` + (or set the environment variables read by :meth:`from_env`) and the + library records spans for session lifecycle steps, every gRPC call at + both sides of the local proxy, the proxy's own connections and the + client-side plan/result conversions. Records go to ``sinks``; + ``attributes`` are attached to every record. + + With ``wrap_client=False`` the pyspark client is left untouched: only + the session, proxy and gRPC spans are recorded. + """ + + def __init__( + self, + sinks: Sequence[Sink] = (), + attributes: Optional[Mapping[str, Any]] = None, + *, + wrap_client: bool = True, + ): + self.sinks: List[Sink] = list(sinks) + self.wrap_client = wrap_client + self._attributes: Dict[str, Any] = dict(attributes or {}) + self._lock = threading.Lock() + self._failed_sinks: set = set() + self.recorder = Recorder(self._emit, self._attribute_snapshot) + self._emit( + {"type": "meta", "wall_ns": time.time_ns(), "attrs": _meta_attrs()} + ) + + @classmethod + def from_env(cls) -> Optional["Instrumentation"]: + """An instance configured from the environment, or ``None``. + + ``MANAGED_SPARK_CONNECT_INSTRUMENTATION_FILE=`` appends JSONL + records to ````; ``MANAGED_SPARK_CONNECT_INSTRUMENTATION_LOG`` + set to a true value logs records at DEBUG level. Both may be set. + """ + sinks: List[Sink] = [] + path = os.environ.get(ENV_FILE) + if path: + sinks.append(JsonlFileSink(path, append=True)) + if os.environ.get(ENV_LOG, "").strip().lower() in _TRUE_VALUES: + sinks.append(LoggingSink()) + if not sinks: + return None + return cls(sinks) + + @property + def enabled(self) -> bool: + return True + + @property + def attributes(self) -> Dict[str, Any]: + """A copy of the global attributes.""" + return self._attribute_snapshot() + + def _attribute_snapshot(self) -> Dict[str, Any]: + with self._lock: + return dict(self._attributes) + + def set_attribute(self, key: str, value: Any) -> None: + """Attach ``key=value`` to every record emitted from now on.""" + with self._lock: + self._attributes[key] = value + + def remove_attribute(self, key: str) -> None: + with self._lock: + self._attributes.pop(key, None) + + @property + def records(self) -> List[Dict[str, Any]]: + """The records of all :class:`InMemorySink` sinks, concatenated.""" + records: List[Dict[str, Any]] = [] + for sink in self.sinks: + if isinstance(sink, InMemorySink): + records.extend(sink.records) + return records + + def wrap_channel(self, channel, layer: str = "outer"): + """``channel`` with this instrumentation's gRPC interceptor added.""" + from ._grpc import wrap_channel + + return wrap_channel(channel, self.recorder, layer) + + def wait_idle(self, timeout_s: float, prefix: str = "grpc.") -> bool: + """Wait for in-flight spans (RPCs by default); see + :meth:`Recorder.wait_idle`.""" + return self.recorder.wait_idle(timeout_s, prefix) + + def flush(self) -> None: + for sink in self.sinks: + self._call_sink(sink, "flush") + + def close(self) -> None: + """Flush and close every sink. Safe to call more than once.""" + for sink in self.sinks: + self._call_sink(sink, "flush") + self._call_sink(sink, "close") + + def _emit(self, record: Dict[str, Any]) -> None: + with self._lock: + if self._attributes and record["type"] != "span": + # Spans carry the globals they started with (see Span); + # events, counters and meta get them as of now. + attrs = record["attrs"] + for key, value in self._attributes.items(): + attrs.setdefault(key, value) + for sink in self.sinks: + try: + sink.write(record) + except Exception as e: + self._warn_sink(sink, "write", e) + + def _call_sink(self, sink: Sink, method: str) -> None: + try: + getattr(sink, method)() + except Exception as e: + self._warn_sink(sink, method, e) + + def _warn_sink(self, sink: Sink, method: str, exc: Exception) -> None: + # Telemetry must never take the user's job down; warn once per sink. + if id(sink) not in self._failed_sinks: + self._failed_sinks.add(id(sink)) + logger.warning( + "Instrumentation sink %r failed in %s(): %s", sink, method, exc + ) + + def __repr__(self) -> str: + return f"Instrumentation(sinks={self.sinks!r})" + + +def get_recorder(instrumentation: Optional[Instrumentation]): + """The recorder of ``instrumentation``, or the no-op recorder.""" + if instrumentation is None: + return NOOP_RECORDER + return instrumentation.recorder diff --git a/google/cloud/managed_spark_connect/instrumentation/_grpc.py b/google/cloud/managed_spark_connect/instrumentation/_grpc.py new file mode 100644 index 00000000..e1941788 --- /dev/null +++ b/google/cloud/managed_spark_connect/instrumentation/_grpc.py @@ -0,0 +1,421 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""A gRPC client interceptor that records one span per RPC. + +The library installs it at two layers of the channel stack: + +* ``outer``: around the whole :class:`ProxiedChannel`, so the span covers + everything the pyspark client waits for, including the session liveness + guard that runs before each RPC (its share is reported as ``guard_ns``). +* ``inner``: around the plain gRPC channel to the local proxy, so the span + is the pure wire time of the RPC. + +Inner spans are children of outer spans; both are children of whatever span +is current on the calling thread (e.g. ``client.to_pandas``). + +Interceptors receive requests before serialization, so ``request_bytes``, +``session_id`` and ``operation_id`` come straight from the request proto. +Streaming responses are counted by a thin iterator proxy that forwards the +``grpc.Call`` surface (``code()``, ``trailing_metadata()``, ``cancel()``, +...) to the real call, because pyspark reads it on errors. +""" + +import contextvars +import functools +import threading +import time +import weakref +from typing import Any, Callable, Dict, Iterator, Optional + +import grpc + +from ._core import _CURRENT_SPAN, NOOP_RECORDER, Recorder, Span + +OUTER = "outer" +INNER = "inner" + +_GUARD_NS: contextvars.ContextVar[int] = contextvars.ContextVar( + "managed_spark_connect_guard_ns", default=0 +) + + +def timed_guard(callback: Optional[Callable[[], bool]]): + """Wrap a session liveness callback so its duration is attributed. + + The time spent in ``callback`` is added to a context variable that the + outer interceptor reads after the RPC and records as ``guard_ns``. + ``None`` stays ``None``; a callback that is already wrapped is returned + as is, so wrapping twice does not double count. + """ + if callback is None or getattr(callback, "_timed_guard", None) is True: + return callback + + @functools.wraps(callback) + def guarded() -> bool: + start = time.perf_counter_ns() + try: + return callback() + finally: + _GUARD_NS.set(_GUARD_NS.get() + time.perf_counter_ns() - start) + + guarded._timed_guard = True # type: ignore[attr-defined] + return guarded + + +def _byte_size(message: Any) -> int: + byte_size = getattr(message, "ByteSize", None) + if byte_size is None: + return 0 + try: + return byte_size() + except Exception: # pragma: no cover - defensive + return 0 + + +def _request_attrs(request: Any) -> Dict[str, Any]: + attrs: Dict[str, Any] = {} + for field in ("session_id", "operation_id"): + value = getattr(request, field, None) + if value: + attrs[field] = value + return attrs + + +def _code_name(call: Any) -> Optional[str]: + code_fn = getattr(call, "code", None) + if code_fn is None: + return None + try: + code = code_fn() + except Exception: + return None + return getattr(code, "name", str(code)) if code is not None else None + + +def _is_pending_future(outcome: Any) -> bool: + done = getattr(outcome, "done", None) + if done is None: + return False + try: + return not done() + except Exception: # pragma: no cover - defensive + return False + + +def _weak_status_callback(proxy: "_StreamingCall") -> Callable[[], None]: + """A status callback that does not keep ``proxy`` alive. + + gRPC holds registered callbacks until the RPC terminates; a strong + reference would form a cycle (call -> callback -> proxy -> call) and + delay the proxy's finalizer, which is what ends the span when the + consumer drops an undrained stream. + """ + ref = weakref.ref(proxy) + + def on_status() -> None: + live = ref() + if live is not None: + live._record_status() + + return on_status + + +class _CountingRequestIterator: + """Counts the request messages of a client-streaming RPC.""" + + def __init__(self, requests: Any, span: Span): + self._requests = iter(requests) + self._span = span + self.count = 0 + self.bytes = 0 + + def __iter__(self) -> "_CountingRequestIterator": + return self + + def __next__(self) -> Any: + message = next(self._requests) + if self.count == 0: + self._span.set(**_request_attrs(message)) + self.count += 1 + self.bytes += _byte_size(message) + self._span.set(request_count=self.count, request_bytes=self.bytes) + return message + + +class _StreamingCall(grpc.Call): + """Iterator proxy around a server-streaming call. + + Records ``first_response`` when the first message arrives and ``status`` + when gRPC reports the final status (from the channel's own thread, via + ``add_callback``, independent of how fast the consumer iterates). The + span ends when the consumer is done with the stream: at the end of the + iteration, on an error, on ``cancel()``, or when the consumer drops the + iterator without draining it (pyspark stops reading an ``ExecutePlan`` + stream as soon as it sees ``result_complete``), which is caught by + ``__del__``. ``response_count``/``response_bytes`` therefore count the + messages delivered to the consumer. + """ + + # Class-level defaults keep __del__/__getattr__ safe on a half-built + # instance (e.g. when __init__ raised). + _call: Any = None + _finished = True + + def __init__(self, call: Any, span: Span): + self._call = call + self._span = span + self._count = 0 + self._bytes = 0 + self._lock = threading.Lock() + self._finished = False + self._status_recorded = False + self._status_mono_ns: Optional[int] = None + add_callback = getattr(call, "add_callback", None) + if add_callback is None or not add_callback( + _weak_status_callback(self) + ): + # No callback support, or the RPC is already over. + self._record_status() + + # -- iteration ----------------------------------------------------- + + def __iter__(self) -> Iterator[Any]: + return self + + def __next__(self) -> Any: + try: + message = next(self._call) + except StopIteration: + self._finish(None) + raise + except BaseException as e: + self._finish(e) + raise + self._count += 1 + self._bytes += _byte_size(message) + if self._count == 1: + self._span.add_event("first_response") + return message + + def _record_status(self) -> None: + # Runs on the gRPC thread (status callback) or the consumer thread + # (_finish); the lock makes the check-and-set atomic so the status + # event is recorded exactly once. + with self._lock: + if self._status_recorded or self._finished: + return + is_active = getattr(self._call, "is_active", None) + if is_active is not None and is_active(): + # code() would block until the RPC completes; the callback + # that fires on completion records the status then. + return + self._status_recorded = True + self._status_mono_ns = time.perf_counter_ns() + code = _code_name(self._call) + if code is not None: + self._span.set(grpc_code=code) + self._span.add_event("status") + + def _finish(self, exc: Optional[BaseException]) -> None: + self._record_status() + with self._lock: + if self._finished: + return + self._finished = True + inner = self._call + if isinstance(inner, _StreamingCall): + # Nested layers: the inner span closes before the outer one. + inner._finish(exc) + self._span.set(response_count=self._count, response_bytes=self._bytes) + if exc is not None: + self._span.record_exception(exc) + # The RPC is over when its status arrived, even if the consumer + # held on to the iterator for a while (pyspark's reattach iterator + # is released only after its ReleaseExecute has been queued). + self._span.end(end_mono_ns=self._status_mono_ns) + + def __del__(self) -> None: + try: + self._finish(None) + except Exception: # pragma: no cover - never raise from a finalizer + pass + + # -- grpc.Call / grpc.RpcContext ----------------------------------- + + def initial_metadata(self): + return self._call.initial_metadata() + + def trailing_metadata(self): + return self._call.trailing_metadata() + + def code(self): + return self._call.code() + + def details(self): + return self._call.details() + + def debug_error_string(self): + return self._call.debug_error_string() + + def is_active(self) -> bool: + return self._call.is_active() + + def time_remaining(self): + return self._call.time_remaining() + + def cancel(self): + result = self._call.cancel() + self._span.set(cancelled=True) + self._finish(None) + return result + + def add_callback(self, callback) -> bool: + return self._call.add_callback(callback) + + def __getattr__(self, name: str) -> Any: + # Anything else (result(), exception(), running(), ...) belongs to + # the real call. + return getattr(self._call, name) + + def __repr__(self) -> str: + return f"_StreamingCall({self._span!r})" + + +class GrpcInterceptor( + grpc.UnaryUnaryClientInterceptor, + grpc.UnaryStreamClientInterceptor, + grpc.StreamUnaryClientInterceptor, + grpc.StreamStreamClientInterceptor, +): + """Records a ``grpc.`` span for every RPC on a channel.""" + + def __init__(self, recorder: Recorder, layer: str = OUTER): + if layer not in (OUTER, INNER): + raise ValueError(f"layer must be 'outer' or 'inner', got {layer!r}") + self._recorder = recorder + self._layer = layer + + # -- helpers ------------------------------------------------------- + + def _start(self, details: grpc.ClientCallDetails, kind: str) -> Span: + method = details.method + if isinstance(method, bytes): + method = method.decode() + return self._recorder.span( + "grpc." + method.rsplit("/", 1)[-1], + method=method, + kind=kind, + layer=self._layer, + ) + + def _run(self, span: Span, continuation: Callable, details, request): + """Call ``continuation`` with ``span`` current, attributing guard time. + + Exceptions raised synchronously (e.g. a failed liveness guard) end + the span as an error and propagate. + """ + span_token = _CURRENT_SPAN.set(span) + guard_token = _GUARD_NS.set(0) if self._layer == OUTER else None + try: + return continuation(details, request) + except BaseException as e: + if guard_token is not None: + span.set(guard_ns=_GUARD_NS.get()) + span.record_exception(e) + span.end() + raise + finally: + if guard_token is not None: + if not span.ended: + span.set(guard_ns=_GUARD_NS.get()) + _GUARD_NS.reset(guard_token) + try: + _CURRENT_SPAN.reset(span_token) + except ValueError: # pragma: no cover - different context + _CURRENT_SPAN.set(None) + + def _finish_unary(self, span: Span, outcome: Any) -> None: + code = _code_name(outcome) + if code is not None: + span.set(grpc_code=code) + span.add_event("status") + try: + response = outcome.result() + except BaseException as e: + span.record_exception(e) + else: + span.set(response_count=1, response_bytes=_byte_size(response)) + span.end() + + def _unary_response(self, span: Span, outcome: Any) -> Any: + if _is_pending_future(outcome): + outcome.add_done_callback( + lambda done: self._finish_unary(span, done) + ) + else: + self._finish_unary(span, outcome) + return outcome + + # -- grpc interceptor interface ----------------------------------- + + def intercept_unary_unary(self, continuation, client_call_details, request): + span = self._start(client_call_details, "unary_unary") + span.set(request_bytes=_byte_size(request), **_request_attrs(request)) + outcome = self._run(span, continuation, client_call_details, request) + return self._unary_response(span, outcome) + + def intercept_unary_stream( + self, continuation, client_call_details, request + ): + span = self._start(client_call_details, "unary_stream") + span.set(request_bytes=_byte_size(request), **_request_attrs(request)) + call = self._run(span, continuation, client_call_details, request) + return _StreamingCall(call, span) + + def intercept_stream_unary( + self, continuation, client_call_details, request_iterator + ): + span = self._start(client_call_details, "stream_unary") + requests = _CountingRequestIterator(request_iterator, span) + outcome = self._run(span, continuation, client_call_details, requests) + return self._unary_response(span, outcome) + + def intercept_stream_stream( + self, continuation, client_call_details, request_iterator + ): + span = self._start(client_call_details, "stream_stream") + requests = _CountingRequestIterator(request_iterator, span) + call = self._run(span, continuation, client_call_details, requests) + return _StreamingCall(call, span) + + +def _as_recorder(instrumentation_or_recorder: Any): + if instrumentation_or_recorder is None: + return NOOP_RECORDER + recorder = getattr(instrumentation_or_recorder, "recorder", None) + return recorder if recorder is not None else instrumentation_or_recorder + + +def wrap_channel( + channel: grpc.Channel, instrumentation_or_recorder: Any, layer: str = OUTER +) -> grpc.Channel: + """``channel`` with a :class:`GrpcInterceptor` at ``layer`` in front. + + Accepts an :class:`Instrumentation`, a :class:`Recorder` or ``None``; + returns ``channel`` itself when instrumentation is disabled. + """ + recorder = _as_recorder(instrumentation_or_recorder) + if not recorder.enabled: + return channel + return grpc.intercept_channel(channel, GrpcInterceptor(recorder, layer)) diff --git a/google/cloud/managed_spark_connect/instrumentation/_sinks.py b/google/cloud/managed_spark_connect/instrumentation/_sinks.py new file mode 100644 index 00000000..89f59696 --- /dev/null +++ b/google/cloud/managed_spark_connect/instrumentation/_sinks.py @@ -0,0 +1,154 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Sinks: where instrumentation records go. + +A sink receives one ``dict`` per record (see the module docstring of +:mod:`google.cloud.managed_spark_connect.instrumentation` for the schema). +Sinks must tolerate being called from several threads at once; the +:class:`Instrumentation` serializes ``write`` calls, but ``flush`` and +``close`` may race with them. +""" + +import json +import logging +import os +import threading +from typing import Any, Callable, Dict, List, Optional, Protocol + + +class Sink(Protocol): + """The sink interface. Any object with these three methods will do.""" + + def write(self, record: Dict[str, Any]) -> None: + ... + + def flush(self) -> None: + ... + + def close(self) -> None: + ... + + +def _to_json(record: Dict[str, Any]) -> str: + # default=str: attribute values are usually scalars, but keep odd ones + # (enums, timestamps) from taking the whole record down. + return json.dumps( + record, separators=(",", ":"), ensure_ascii=False, default=str + ) + + +class JsonlFileSink: + """Appends one JSON object per line to a file. + + The file is line buffered, so every record reaches the OS as soon as it + is written and a crash loses at most the record being written. A closed + sink reopens its file (in append mode) on the next ``write``, so one sink + can outlive the session that closed it. + """ + + def __init__(self, path: "str | os.PathLike[str]", *, append: bool = True): + self.path = os.fspath(path) + self._append = append + self._lock = threading.Lock() + self._file = None + self._open(truncate=not append) + + def _open(self, truncate: bool) -> None: + self._file = open( + self.path, + "w" if truncate else "a", + buffering=1, + encoding="utf-8", + newline="\n", + ) + + def write(self, record: Dict[str, Any]) -> None: + line = _to_json(record) + "\n" + with self._lock: + if self._file is None: + self._open(truncate=False) + self._file.write(line) + + def flush(self) -> None: + with self._lock: + if self._file is not None: + self._file.flush() + + def close(self) -> None: + with self._lock: + if self._file is not None: + self._file.close() + self._file = None + + def __repr__(self) -> str: + return f"JsonlFileSink({self.path!r})" + + +class InMemorySink: + """Keeps records in a list, for tests and short interactive sessions.""" + + def __init__(self): + self.records: List[Dict[str, Any]] = [] + + def write(self, record: Dict[str, Any]) -> None: + self.records.append(record) + + def flush(self) -> None: + pass + + def close(self) -> None: + pass + + def clear(self) -> None: + self.records.clear() + + +class LoggingSink: + """Logs every record as one JSON line through :mod:`logging`.""" + + def __init__( + self, + logger: Optional[logging.Logger] = None, + level: int = logging.DEBUG, + ): + self.logger = logger or logging.getLogger( + "google.cloud.managed_spark_connect.instrumentation" + ) + self.level = level + + def write(self, record: Dict[str, Any]) -> None: + if self.logger.isEnabledFor(self.level): + self.logger.log(self.level, "%s", _to_json(record)) + + def flush(self) -> None: + pass + + def close(self) -> None: + pass + + +class CallbackSink: + """Hands every record to a callable, for custom exporters.""" + + def __init__(self, fn: Callable[[Dict[str, Any]], None]): + self.fn = fn + + def write(self, record: Dict[str, Any]) -> None: + self.fn(record) + + def flush(self) -> None: + pass + + def close(self) -> None: + pass diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index d6470acd..6b70a953 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -45,6 +45,13 @@ from google.auth.exceptions import DefaultCredentialsError from google.cloud.managed_spark_connect.client import ManagedSparkChannelBuilder from google.cloud.managed_spark_connect.exceptions import ManagedSparkConnectException +from google.cloud.managed_spark_connect.instrumentation import ( + NOOP_RECORDER, + Instrumentation, + get_recorder, + install_client_wrappers, + uninstall_client_wrappers, +) from google.cloud.managed_spark_connect.pypi_artifacts import PyPiArtifacts from google.cloud.dataproc_v1 import ( AuthenticationConfig, @@ -267,6 +274,39 @@ def _is_valid_session_id(session_id: str) -> bool: return bool(re.match(pattern, session_id)) +def _timestamp_str(value: Any) -> Optional[str]: + if value is None: + return None + isoformat = getattr(value, "isoformat", None) + return isoformat() if callable(isoformat) else str(value) + + +def _session_state_attrs(session: Any) -> Dict[str, Any]: + """Instrumentation attributes describing a Dataproc session's lifecycle. + + Best effort: anything missing (e.g. on a mocked session) is left out. + """ + attrs: Dict[str, Any] = {} + try: + attrs["create_time"] = _timestamp_str(session.create_time) + attrs["state_time"] = _timestamp_str(session.state_time) + state = getattr(session, "state", None) + attrs["state"] = getattr(state, "name", str(state)) + attrs["state_history"] = json.dumps( + [ + { + "state": getattr(h.state, "name", str(h.state)), + "state_message": h.state_message, + "state_start_time": _timestamp_str(h.state_start_time), + } + for h in session.state_history + ] + ) + except Exception: # pragma: no cover - mocks, partial protos + pass + return attrs + + class ManagedSparkSession(SparkSession): """The entry point to programming Spark with the Dataset and DataFrame API. @@ -297,6 +337,9 @@ class ManagedSparkSession(SparkSession): _active_session_uses_custom_id: ClassVar[bool] = False _execution_progress_bar = dict() + #: The :class:`Instrumentation` this session records to, or ``None``. + instrumentation: Optional[Instrumentation] = None + class Builder(SparkSession.Builder): def __init__(self): @@ -316,6 +359,18 @@ def __init__(self): SessionControllerClient ] = None self._session_template_controller_client = None + self._instrumentation: Optional[Instrumentation] = None + self._last_session_response: Optional[Session] = None + + @property + def _recorder(self): + return get_recorder(self._instrumentation) + + def instrumentation(self, instrumentation: Optional[Instrumentation]): + """Record what the library does (see the ``instrumentation`` + package). Takes precedence over the environment variables.""" + self._instrumentation = instrumentation + return self @property def session_controller_client(self) -> SessionControllerClient: @@ -471,6 +526,13 @@ def create(self) -> "ManagedSparkSession": def __create_spark_connect_session_from_s8s( self, session_response, session_name ) -> "ManagedSparkSession": + recorder = self._recorder + recorder.trace_id = session_response.uuid + self._last_session_response = session_response + with recorder.span("session.connect"): + return self.__connect(session_response, session_name) + + def __connect(self, session_response, session_name): ManagedSparkSession._active_s8s_session_uuid = session_response.uuid ManagedSparkSession._project_id = self._project_id ManagedSparkSession._region = self._region @@ -480,15 +542,27 @@ def __create_spark_connect_session_from_s8s( ) url = f"{spark_connect_url}/;session_id={session_response.uuid};use_ssl=true" logger.debug(f"Spark Connect URL: {url}") + recorder = self._recorder + + def is_active(): + with recorder.span( + "dataproc.get_session", reason="rpc_guard", new_client=True + ): + return is_s8s_session_active( + session_name, self._client_options + ) + self._channel_builder = ManagedSparkChannelBuilder( url, - is_active_callback=lambda: is_s8s_session_active( - session_name, self._client_options - ), + is_active_callback=is_active, + instrumentation=self._instrumentation, ) assert self._channel_builder is not None - session = ManagedSparkSession(connection=self._channel_builder) + session = ManagedSparkSession( + connection=self._channel_builder, + instrumentation=self._instrumentation, + ) # Register handler for Cell Execution Progress bar session._register_progress_execution_handler() @@ -586,21 +660,26 @@ def create_session_pbar(): self._client_options, ) ) - operation = SessionControllerClient( - client_options=self._client_options - ).create_session(session_request) + with self._recorder.span("session.create_rpc"): + operation = SessionControllerClient( + client_options=self._client_options + ).create_session(session_request) self._display_session_link_on_creation(session_id) self._display_view_session_details_button(session_id) create_session_pbar_thread.start() - session_response: Session = operation.result( - polling=retry.Retry( - predicate=POLLING_PREDICATE, - initial=5.0, # seconds - maximum=5.0, # seconds - multiplier=1.0, - timeout=1200, # seconds + with self._recorder.span( + "session.wait_active", poll_interval_s=5.0 + ) as wait_span: + session_response: Session = operation.result( + polling=retry.Retry( + predicate=POLLING_PREDICATE, + initial=5.0, # seconds + maximum=5.0, # seconds + multiplier=1.0, + timeout=1200, # seconds + ) ) - ) + wait_span.set(**_session_state_attrs(session_response)) stop_create_session_pbar_event.set() create_session_pbar_thread.join() file_path = ( @@ -666,13 +745,26 @@ def create_session_pbar(): def _wait_for_session_available( self, session_name: str, timeout: int = 1200 + ) -> Session: + with self._recorder.span( + "session.wait_endpoint", poll_interval_s=5.0 + ) as span: + session = self._poll_for_session_endpoint(session_name, timeout) + span.set(**_session_state_attrs(session)) + return session + + def _poll_for_session_endpoint( + self, session_name: str, timeout: int ) -> Session: start_time = time.time() while time.time() - start_time < timeout: try: - session = self.session_controller_client.get_session( - name=session_name - ) + with self._recorder.span( + "dataproc.get_session", reason="poll", new_client=False + ): + session = self.session_controller_client.get_session( + name=session_name + ) if "Spark Connect Server" in session.runtime_info.endpoints: return session if session.state in ( @@ -770,45 +862,63 @@ def _get_exiting_active_session( session_name = f"projects/{self._project_id}/locations/{self._region}/sessions/{s8s_session_id}" session_response = None session = None + recorder = self._recorder if s8s_session_id is not None: - session_response = get_active_s8s_session_response( - session_name, self._client_options + reuse_span = recorder.span( + "session.reuse", session_id=s8s_session_id ) - session = ManagedSparkSession.getActiveSession() - - if session is None: - session = ManagedSparkSession._default_session + else: + reuse_span = NOOP_RECORDER.span("session.reuse") + with reuse_span: + if s8s_session_id is not None: + with recorder.span( + "dataproc.get_session", reason="reuse", new_client=True + ): + session_response = get_active_s8s_session_response( + session_name, self._client_options + ) + if session_response is not None: + recorder.trace_id = session_response.uuid + reuse_span.set(found=session_response is not None) + session = ManagedSparkSession.getActiveSession() - if session_response is not None: if session is None: - session_response = self._wait_for_session_available( - session_name - ) - session = self.__create_spark_connect_session_from_s8s( - session_response, session_name + session = ManagedSparkSession._default_session + + if session_response is not None: + self._last_session_response = session_response + if session is None: + session_response = self._wait_for_session_available( + session_name + ) + session = self.__create_spark_connect_session_from_s8s( + session_response, session_name + ) + links = _session_links( + self._region, + self._project_id, + s8s_session_id, + html=False, ) - links = _session_links( - self._region, self._project_id, s8s_session_id, html=False - ) - versions = self._session_versions(session) - subject = ( - f"Using existing Managed Spark Session {s8s_session_id} " - f"{versions}" - ).rstrip() - print( - f"{subject}; configuration changes may not be applied: " - f"{links.url}" - ) - self._display_view_session_details_button(s8s_session_id) - return session - else: - if session is not None: + versions = self._session_versions(session) + subject = ( + f"Using existing Managed Spark Session {s8s_session_id} " + f"{versions}" + ).rstrip() print( - f"{s8s_session_id} Managed Spark Session is not active, stopping and creating a new one" + f"{subject}; configuration changes may not be applied: " + f"{links.url}" ) - session.stop() + self._display_view_session_details_button(s8s_session_id) + return session + else: + if session is not None: + print( + f"{s8s_session_id} Managed Spark Session is not active, stopping and creating a new one" + ) + session.stop() - return None + return None def getOrCreate(self) -> "ManagedSparkSession": with ManagedSparkSession._lock: @@ -829,13 +939,35 @@ def getOrCreate(self) -> "ManagedSparkSession": f"Error while creating Managed Spark Session: location is not set" ) - # Handle custom session ID by setting it early and letting existing logic handle it - if self._custom_session_id: - self._handle_custom_session_id() - - session = self._get_exiting_active_session() - if session is None: - session = self.__create() + if self._instrumentation is None: + existing_session = ( + ManagedSparkSession.getActiveSession() + or ManagedSparkSession._default_session + ) + if ( + existing_session is not None + and existing_session.instrumentation is not None + and ( + self._custom_session_id is None + or self._custom_session_id + == ManagedSparkSession._active_s8s_session_id + ) + ): + self._instrumentation = existing_session.instrumentation + else: + self._instrumentation = Instrumentation.from_env() + + with self._recorder.span("session.get_or_create") as span: + # Handle custom session ID by setting it early and letting existing logic handle it + if self._custom_session_id: + self._handle_custom_session_id() + + session = self._get_exiting_active_session() + mode = "reuse" + if session is None: + mode = "create" + session = self.__create() + span.set(mode=mode, **self._session_span_attrs()) # Register this session as the instantiated SparkSession for compatibility # with tools and libraries that expect SparkSession._instantiatedSession @@ -845,6 +977,21 @@ def getOrCreate(self) -> "ManagedSparkSession": return session + def _session_span_attrs(self) -> Dict[str, Any]: + """Identity attributes of the session for ``session.*`` spans.""" + attrs: Dict[str, Any] = { + "session_id": ManagedSparkSession._active_s8s_session_id, + "uuid": ManagedSparkSession._active_s8s_session_uuid, + } + response = self._last_session_response + if response is not None: + try: + attrs["runtime_version"] = response.runtime_config.version + attrs["template"] = response.session_template + except Exception: # pragma: no cover - mocks + pass + return attrs + def _handle_custom_session_id(self): """Handle custom session ID by checking if it exists and setting _active_s8s_session_id.""" session_response = self._get_session_by_id(self._custom_session_id) @@ -1160,9 +1307,12 @@ def _get_session_by_id(self, session_id: str) -> Optional[Session]: try: get_request = GetSessionRequest(name=session_name) - session = self.session_controller_client.get_session( - get_request - ) + with self._recorder.span( + "dataproc.get_session", reason="reuse", new_client=False + ): + session = self.session_controller_client.get_session( + get_request + ) logger.debug( f"Found existing session {session_id} in state: {session.state}" @@ -1217,9 +1367,12 @@ def _wait_for_termination(self, session_name: str, timeout: int = 180): while time.time() - start_time < timeout: try: get_request = GetSessionRequest(name=session_name) - session = self.session_controller_client.get_session( - get_request - ) + with self._recorder.span( + "dataproc.get_session", reason="stop", new_client=False + ): + session = self.session_controller_client.get_session( + get_request + ) if session.state in [ Session.State.TERMINATED, @@ -1257,6 +1410,7 @@ def __init__( self, connection: Union[str, ManagedSparkChannelBuilder], user_id: Optional[str] = None, + instrumentation: Optional[Instrumentation] = None, ): """ Creates a new ManagedSparkSession for the Spark Connect interface. @@ -1270,10 +1424,17 @@ def __init__( user_id : str, optional If not set, will default to the $USER environment. Defining the user ID as part of the connection string takes precedence. + instrumentation : :class:`Instrumentation`, optional + Records spans for the client-side work of this session (see the + ``instrumentation`` package). Exposed as ``self.instrumentation``. """ super().__init__(connection, user_id) + self.instrumentation = instrumentation + if instrumentation is not None and instrumentation.wrap_client: + install_client_wrappers(self.client, instrumentation.recorder) + execute_plan_request_base_method = ( self.client._execute_plan_request_with_metadata ) @@ -1594,6 +1755,17 @@ def stop(self, terminate: Optional[bool] = None) -> None: >>> spark.stop(terminate=False) """ + recorder = get_recorder(self.instrumentation) + try: + with recorder.span("session.stop") as span: + self._stop(terminate, span) + finally: + if self.instrumentation is not None: + uninstall_client_wrappers(self.client) + recorder.trace_id = "" + self.instrumentation.close() + + def _stop(self, terminate: Optional[bool], span) -> None: with ManagedSparkSession._lock: if ManagedSparkSession._active_s8s_session_id is not None: # Determine if we should terminate the server-side session @@ -1605,6 +1777,10 @@ def stop(self, terminate: Optional[bool] = None) -> None: else: should_terminate = terminate + span.set( + terminate=should_terminate, + session_id=ManagedSparkSession._active_s8s_session_id, + ) if should_terminate: # Terminate the server-side session logger.debug( @@ -1615,6 +1791,7 @@ def stop(self, terminate: Optional[bool] = None) -> None: ManagedSparkSession._region, ManagedSparkSession._active_s8s_session_id, self._client_options, + recorder=get_recorder(self.instrumentation), ) else: # Client-side cleanup only @@ -1654,10 +1831,16 @@ def stop(self, terminate: Optional[bool] = None) -> None: def terminate_s8s_session( - project_id, region, active_s8s_session_id, client_options=None + project_id, + region, + active_s8s_session_id, + client_options=None, + recorder=None, ): from google.cloud.dataproc_v1 import SessionControllerClient + if recorder is None: + recorder = NOOP_RECORDER logger.debug(f"Terminating Managed Spark Session: {active_s8s_session_id}") terminate_session_request = TerminateSessionRequest() session_name = f"projects/{project_id}/locations/{region}/sessions/{active_s8s_session_id}" @@ -1665,18 +1848,26 @@ def terminate_s8s_session( state = None try: session_client = SessionControllerClient(client_options=client_options) - session_client.terminate_session(terminate_session_request) + with recorder.span("session.terminate_rpc"): + session_client.terminate_session(terminate_session_request) get_session_request = GetSessionRequest() get_session_request.name = session_name state = Session.State.ACTIVE - while ( - state != Session.State.TERMINATING - and state != Session.State.TERMINATED - and state != Session.State.FAILED - ): - session = session_client.get_session(get_session_request) - state = session.state - time.sleep(1) + with recorder.span( + "session.wait_terminating", poll_interval_s=1.0 + ) as wait_span: + while ( + state != Session.State.TERMINATING + and state != Session.State.TERMINATED + and state != Session.State.FAILED + ): + with recorder.span( + "dataproc.get_session", reason="stop", new_client=False + ): + session = session_client.get_session(get_session_request) + state = session.state + time.sleep(1) + wait_span.set(state=getattr(state, "name", str(state))) except NotFound: logger.debug( f"{active_s8s_session_id} Managed Spark Session already deleted" diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py new file mode 100644 index 00000000..7d3869e7 --- /dev/null +++ b/tests/unit/test_core.py @@ -0,0 +1,185 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the channel builder and the guarded proxied channel. + +No network: the local proxy and the real gRPC channel are replaced by fakes. +""" + +from unittest import mock + +import grpc +import pytest + +from google.cloud.managed_spark_connect.client import core + + +class FakeMultiCallable: + """Stands in for the object returned by ``grpc.Channel.unary_unary``.""" + + def __init__(self): + self.calls = [] + + def __call__(self, request, **kwargs): + self.calls.append(("call", request)) + return f"response:{request}" + + def with_call(self, request, **kwargs): + self.calls.append(("with_call", request)) + return f"response:{request}", mock.Mock(spec=grpc.Call) + + def future(self, request, **kwargs): + self.calls.append(("future", request)) + return mock.Mock(spec=grpc.Future) + + +class FakeChannel(grpc.Channel): + """A grpc.Channel whose multi-callables are :class:`FakeMultiCallable`.""" + + def __init__(self): + self.multicallables = {} + self.closed = False + + def _get(self, method): + return self.multicallables.setdefault(method, FakeMultiCallable()) + + def unary_unary(self, method, *args, **kwargs): + return self._get(method) + + def unary_stream(self, method, *args, **kwargs): + return self._get(method) + + def stream_unary(self, method, *args, **kwargs): + return self._get(method) + + def stream_stream(self, method, *args, **kwargs): + return self._get(method) + + def subscribe(self, callback, try_to_connect=False): + return None + + def unsubscribe(self, callback): + return None + + def close(self): + self.closed = True + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() + return False + + +@pytest.fixture +def fake_channel(): + channel = FakeChannel() + with ( + mock.patch.object(core.proxy, "ManagedSparkSessionProxy") as proxy_cls, + mock.patch.object(core, "DefaultChannelBuilder") as builder_cls, + ): + proxy_cls.return_value.port = 12345 + builder_cls.return_value.toChannel.return_value = channel + yield channel + + +def test_guarded_multicallable_exposes_with_call_and_future(fake_channel): + guard = mock.Mock(return_value=True) + channel = core.ProxiedChannel("example.com", guard) + + method = channel.unary_unary("/svc/Method") + assert method("a") == "response:a" + response, call = method.with_call("b") + assert response == "response:b" + assert isinstance(call, grpc.Call) + assert method.future("c") is not None + + assert guard.call_count == 3 + assert fake_channel.multicallables["/svc/Method"].calls == [ + ("call", "a"), + ("with_call", "b"), + ("future", "c"), + ] + + +def test_guard_failure_raises_before_the_rpc(fake_channel): + channel = core.ProxiedChannel("example.com", lambda: False) + method = channel.unary_stream("/svc/Stream") + with pytest.raises(RuntimeError, match="Session not active"): + method("a") + with pytest.raises(RuntimeError, match="Session not active"): + method.with_call("a") + assert fake_channel.multicallables["/svc/Stream"].calls == [] + + +def test_no_guard_returns_the_multicallable_unchanged(fake_channel): + channel = core.ProxiedChannel("example.com", None) + assert ( + channel.unary_unary("/svc/Method") + is fake_channel.multicallables["/svc/Method"] + ) + # subscribe/unsubscribe are forwarded, not wrapped. + assert channel.subscribe(lambda *_: None) is None + assert channel.unsubscribe(lambda *_: None) is None + + +def test_proxied_channel_works_under_intercept_channel(fake_channel): + """grpc.intercept_channel calls with_call() on unary-unary methods.""" + + class Recording(grpc.UnaryUnaryClientInterceptor): + + def __init__(self): + self.methods = [] + + def intercept_unary_unary(self, continuation, details, request): + self.methods.append(details.method) + return continuation(details, request) + + guard = mock.Mock(return_value=True) + interceptor = Recording() + channel = grpc.intercept_channel( + core.ProxiedChannel("example.com", guard), interceptor + ) + assert channel.unary_unary("/svc/Method")("req") == "response:req" + assert interceptor.methods == ["/svc/Method"] + assert guard.call_count == 1 + + +def test_channel_builder_honors_add_interceptor(fake_channel): + + class Recording(grpc.UnaryUnaryClientInterceptor): + + def __init__(self): + self.methods = [] + + def intercept_unary_unary(self, continuation, details, request): + self.methods.append(details.method) + return continuation(details, request) + + interceptor = Recording() + builder = core.ManagedSparkChannelBuilder( + "sc://example.com:443/;use_ssl=true", is_active_callback=lambda: True + ) + builder.add_interceptor(interceptor) + channel = builder.toChannel() + + assert channel.unary_unary("/svc/Method")("req") == "response:req" + assert interceptor.methods == ["/svc/Method"] + + +def test_channel_builder_without_interceptors_returns_proxied_channel( + fake_channel, +): + builder = core.ManagedSparkChannelBuilder("sc://example.com:443") + assert isinstance(builder.toChannel(), core.ProxiedChannel) diff --git a/tests/unit/test_instrumentation.py b/tests/unit/test_instrumentation.py new file mode 100644 index 00000000..a86ec52a --- /dev/null +++ b/tests/unit/test_instrumentation.py @@ -0,0 +1,1525 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for google.cloud.managed_spark_connect.instrumentation (no network).""" + +import json +import logging +import os +import threading +from unittest import mock + +import grpc +import pytest +from pyspark.sql.connect import proto + +from google.cloud.managed_spark_connect import instrumentation as inst_lib +from google.cloud.managed_spark_connect.client import core +from google.cloud.managed_spark_connect.instrumentation import ( + CallbackSink, + GrpcInterceptor, + InMemorySink, + Instrumentation, + JsonlFileSink, + LoggingSink, + NOOP_RECORDER, + Recorder, + get_recorder, + install_client_wrappers, + timed_guard, + uninstall_client_wrappers, + wrap_channel, +) +from google.cloud.managed_spark_connect.instrumentation import _client + +SPAN_KEYS = { + "type", + "name", + "span_id", + "parent_id", + "trace_id", + "start_wall_ns", + "start_mono_ns", + "duration_ns", + "status", + "error", + "thread", + "attrs", + "events", +} + + +def spans(inst, name=None): + return [ + r + for r in inst.records + if r["type"] == "span" and (name is None or r["name"] == name) + ] + + +def one(inst, name): + found = spans(inst, name) + assert len(found) == 1, f"expected one {name!r} span, got {found}" + return found[0] + + +@pytest.fixture +def inst(): + return Instrumentation([InMemorySink()], {"run_id": "r1"}) + + +# -------------------------------------------------------------------------- +# Recorder / Span / Instrumentation +# -------------------------------------------------------------------------- + + +def test_meta_record_is_first_and_carries_versions(inst): + meta = inst.records[0] + assert meta["type"] == "meta" + assert meta["attrs"]["run_id"] == "r1" + for key in ("library_version", "pyspark_version", "grpcio_version"): + assert key in meta["attrs"] + assert meta["attrs"]["python"] + + +def test_span_schema_nesting_and_events(inst): + with inst.recorder.span("outer", a=1) as outer: + with inst.recorder.span("inner") as inner: + inner.add_event("first_response") + inner.set(b=2) + o, i = one(inst, "outer"), one(inst, "inner") + assert set(o) == SPAN_KEYS + assert o["parent_id"] is None + assert i["parent_id"] == o["span_id"] + assert len(o["span_id"]) == 16 and int(o["span_id"], 16) >= 0 + assert o["attrs"] == {"run_id": "r1", "a": 1} + assert i["attrs"] == {"run_id": "r1", "b": 2} + assert i["events"][0]["name"] == "first_response" + assert 0 <= i["events"][0]["offset_ns"] <= i["duration_ns"] + assert o["status"] == "ok" and o["error"] is None + assert o["thread"] == threading.current_thread().name + assert i["start_mono_ns"] >= o["start_mono_ns"] + assert o["duration_ns"] >= i["duration_ns"] + # Records are JSON serializable as emitted. + json.dumps(inst.records) + + +def test_span_records_exception_with_grpc_code(inst): + class FakeRpcError(grpc.RpcError): + + def code(self): + return grpc.StatusCode.UNAVAILABLE + + def __str__(self): + return "gone" + + with pytest.raises(FakeRpcError): + with inst.recorder.span("rpc"): + raise FakeRpcError() + rec = one(inst, "rpc") + assert rec["status"] == "error" + assert rec["error"] == { + "type": "FakeRpcError", + "message": "gone", + "grpc_code": "UNAVAILABLE", + } + + +def test_span_records_api_core_exception_code(inst): + from google.api_core.exceptions import NotFound + + with pytest.raises(NotFound): + with inst.recorder.span("x"): + raise NotFound("nope") + assert one(inst, "x")["error"]["grpc_code"] == "NOT_FOUND" + + +def test_global_attributes_apply_to_subsequent_records(inst): + inst.set_attribute("trial_id", "t1") + inst.recorder.event("e") + inst.remove_attribute("trial_id") + inst.recorder.counter("c", 5, unit="bytes") + events = [r for r in inst.records if r["type"] == "event"] + counters = [r for r in inst.records if r["type"] == "counter"] + assert events[0]["attrs"] == {"run_id": "r1", "trial_id": "t1"} + assert {"name", "wall_ns", "mono_ns", "thread"} <= set(events[0]) + assert counters[0]["delta"] == 5 + assert counters[0]["attrs"] == {"run_id": "r1", "unit": "bytes"} + + +def test_trace_id_from_recorder(inst): + inst.recorder.trace_id = "uuid-1" + with inst.recorder.span("a"): + pass + assert one(inst, "a")["trace_id"] == "uuid-1" + # A span started before the trace id was known picks it up at end. + inst.recorder.trace_id = "" + span = inst.recorder.start_span("b") + inst.recorder.trace_id = "uuid-2" + span.end() + assert one(inst, "b")["trace_id"] == "uuid-2" + + +def test_start_span_on_a_thread_has_no_parent(inst): + def worker(): + span = inst.recorder.start_span("thread-owned", conn=1) + span.end() + + with inst.recorder.span("main"): + t = threading.Thread(target=worker, name="proxy-1") + t.start() + t.join() + explicit = inst.recorder.start_span( + "explicit", parent=inst.recorder.current_span() + ) + explicit.end() + owned = one(inst, "thread-owned") + assert owned["parent_id"] is None + assert owned["thread"] == "proxy-1" + assert owned["attrs"]["conn"] == 1 + assert one(inst, "explicit")["parent_id"] == one(inst, "main")["span_id"] + + +def test_end_is_idempotent(inst): + span = inst.recorder.start_span("once") + span.end() + span.add_event("late") + span.end() + assert len(spans(inst, "once")) == 1 + assert one(inst, "once")["events"] == [] + + +def test_noop_recorder_is_allocation_free_and_disabled(): + assert NOOP_RECORDER.enabled is False + a = NOOP_RECORDER.span("x", k=1) + b = NOOP_RECORDER.span("y") + assert a is b + with a as s: + s.set(z=1) + s.add_event("e") + s.record_exception(ValueError()) + assert NOOP_RECORDER.start_span("z") is a + NOOP_RECORDER.event("e") + NOOP_RECORDER.counter("c") + NOOP_RECORDER.trace_id = "ignored" + assert NOOP_RECORDER.trace_id == "" + assert get_recorder(None) is NOOP_RECORDER + inst = Instrumentation() + assert get_recorder(inst) is inst.recorder + assert inst.records == [] + + +def test_sink_failures_are_contained(inst, caplog): + class Broken: + + def write(self, record): + raise IOError("disk full") + + def flush(self): + pass + + def close(self): + pass + + inst.sinks.append(Broken()) + with caplog.at_level(logging.WARNING): + inst.recorder.event("a") + inst.recorder.event("b") + assert [r["name"] for r in inst.records if r["type"] == "event"] == [ + "a", + "b", + ] + assert sum("disk full" in m for m in caplog.messages) == 1 + + +# -------------------------------------------------------------------------- +# Sinks and from_env +# -------------------------------------------------------------------------- + + +def test_jsonl_file_sink_appends_and_reopens(tmp_path): + path = tmp_path / "events.jsonl" + sink = JsonlFileSink(path) + sink.write({"type": "event", "name": "a", "attrs": {}}) + sink.close() + sink.write({"type": "event", "name": "b", "attrs": {"v": b"x"}}) + sink.close() + sink.close() # idempotent + lines = path.read_text(encoding="utf-8").splitlines() + assert [json.loads(l)["name"] for l in lines] == ["a", "b"] + assert json.loads(lines[1])["attrs"]["v"] == "b'x'" + + JsonlFileSink(path, append=False).close() + assert path.read_text() == "" + + +def test_logging_and_callback_sinks(caplog): + seen = [] + log = logging.getLogger("test.instrumentation") + inst = Instrumentation( + [LoggingSink(log, logging.INFO), CallbackSink(seen.append)] + ) + with caplog.at_level(logging.INFO, logger="test.instrumentation"): + inst.recorder.event("hello") + assert any('"name":"hello"' in m for m in caplog.messages) + assert [r["type"] for r in seen] == ["meta", "event"] + + +def test_from_env(tmp_path, monkeypatch): + monkeypatch.delenv(inst_lib.ENV_FILE, raising=False) + monkeypatch.delenv(inst_lib.ENV_LOG, raising=False) + assert Instrumentation.from_env() is None + + monkeypatch.setenv(inst_lib.ENV_LOG, "1") + inst = Instrumentation.from_env() + assert [type(s) for s in inst.sinks] == [LoggingSink] + + path = tmp_path / "out.jsonl" + monkeypatch.setenv(inst_lib.ENV_FILE, str(path)) + monkeypatch.setenv(inst_lib.ENV_LOG, "false") + inst = Instrumentation.from_env() + assert [type(s) for s in inst.sinks] == [JsonlFileSink] + inst.close() + assert json.loads(path.read_text().splitlines()[0])["type"] == "meta" + + +# -------------------------------------------------------------------------- +# gRPC interceptor +# -------------------------------------------------------------------------- + + +class FakeCall(grpc.Call): + """A finished unary call.""" + + def __init__(self, code=grpc.StatusCode.OK): + self._code = code + + def initial_metadata(self): + return (("i", "1"),) + + def trailing_metadata(self): + return (("t", "1"),) + + def code(self): + return self._code + + def details(self): + return "details" + + def is_active(self): + return False + + def time_remaining(self): + return None + + def cancel(self): + return False + + def add_callback(self, callback): + return False + + +class FakeStreamCall(FakeCall): + """A server-streaming call that yields ``messages`` then ends. + + ``error`` (an RpcError instance) is raised after the messages. Status + callbacks fire when the stream is exhausted, like gRPC's spin thread. + """ + + def __init__(self, messages, error=None): + super().__init__(error.code() if error else grpc.StatusCode.OK) + self._messages = iter(messages) + self._error = error + self._callbacks = [] + self._done = False + self.cancelled = False + + def __iter__(self): + return self + + def __next__(self): + try: + return next(self._messages) + except StopIteration: + self._finish() + if self._error is not None: + raise self._error + raise + + def _finish(self): + self._done = True + callbacks, self._callbacks = self._callbacks, [] + for cb in callbacks: + cb() + + def is_active(self): + return not self._done + + def cancel(self): + self.cancelled = True + self._done = True + return True + + def add_callback(self, callback): + if self._done: + return False + self._callbacks.append(callback) + return True + + def debug_error_string(self): + return "debug" + + +class FakeRpcError(grpc.RpcError, grpc.Call): + + def __init__(self, code): + self._code = code + + def code(self): + return self._code + + def details(self): + return "boom" + + def initial_metadata(self): + return () + + def trailing_metadata(self): + return () + + def is_active(self): + return False + + def time_remaining(self): + return None + + def cancel(self): + return False + + def add_callback(self, callback): + return False + + def result(self, timeout=None): + # Like grpc._channel._InactiveRpcError: the error is its own future. + raise self + + +class FakeMultiCallable: + + def __init__(self, channel, method): + self.channel = channel + self.method = method + + def _respond(self, request): + handler = self.channel.handlers.get(self.method) + return handler(request) if handler else proto.ConfigResponse() + + def __call__(self, request, **kwargs): + result = self._respond(request) + if isinstance(result, Exception): + raise result + return result + + def with_call(self, request, **kwargs): + result = self._respond(request) + if isinstance(result, Exception): + raise result + return result, FakeCall() + + def future(self, request, **kwargs): + raise NotImplementedError + + +class FakeChannel(grpc.Channel): + """A channel whose RPCs are answered by ``handlers[method](request)``.""" + + def __init__(self): + self.handlers = {} + self.closed = False + + def unary_unary(self, method, *args, **kwargs): + return FakeMultiCallable(self, method) + + unary_stream = stream_unary = stream_stream = unary_unary + + def subscribe(self, callback, try_to_connect=False): + pass + + def unsubscribe(self, callback): + pass + + def close(self): + self.closed = True + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() + return False + + +EXECUTE = "/spark.connect.SparkConnectService/ExecutePlan" +CONFIG = "/spark.connect.SparkConnectService/Config" +ARTIFACTS = "/spark.connect.SparkConnectService/AddArtifacts" + + +def execute_request(): + return proto.ExecutePlanRequest(session_id="sess-1", operation_id="op-1") + + +def responses(n): + return [ + proto.ExecutePlanResponse(session_id="sess-1", operation_id="op-1") + for _ in range(n) + ] + + +def test_unary_unary_span(inst): + fake = FakeChannel() + fake.handlers[CONFIG] = lambda req: proto.ConfigResponse(session_id="s") + channel = wrap_channel(fake, inst, "inner") + request = proto.ConfigRequest(session_id="sess-1") + response = channel.unary_unary(CONFIG)(request) + assert response.session_id == "s" + rec = one(inst, "grpc.Config") + assert rec["attrs"] == { + "run_id": "r1", + "method": CONFIG, + "kind": "unary_unary", + "layer": "inner", + "request_bytes": request.ByteSize(), + "session_id": "sess-1", + "grpc_code": "OK", + "response_count": 1, + "response_bytes": response.ByteSize(), + } + assert [e["name"] for e in rec["events"]] == ["status"] + assert rec["status"] == "ok" + + +def test_unary_unary_error(inst): + fake = FakeChannel() + fake.handlers[CONFIG] = lambda req: FakeRpcError( + grpc.StatusCode.UNAVAILABLE + ) + channel = wrap_channel(fake, inst.recorder, "outer") + with pytest.raises(grpc.RpcError): + channel.unary_unary(CONFIG)(proto.ConfigRequest()) + rec = one(inst, "grpc.Config") + assert rec["status"] == "error" + assert rec["error"]["grpc_code"] == "UNAVAILABLE" + assert rec["attrs"]["grpc_code"] == "UNAVAILABLE" + assert rec["attrs"]["guard_ns"] == 0 + + +def test_unary_stream_span_and_iterator_proxy(inst): + fake = FakeChannel() + msgs = responses(3) + fake.handlers[EXECUTE] = lambda req: FakeStreamCall(msgs) + channel = wrap_channel(fake, inst, "outer") + call = channel.unary_stream(EXECUTE)(execute_request()) + + assert isinstance(call, grpc.Call) + assert call.initial_metadata() == (("i", "1"),) + assert call.trailing_metadata() == (("t", "1"),) + assert call.details() == "details" + assert call.debug_error_string() == "debug" + assert call.time_remaining() is None + assert call.is_active() is True + assert spans(inst, "grpc.ExecutePlan") == [] # not ended yet + + got = list(call) + assert got == msgs + assert call.code() == grpc.StatusCode.OK + rec = one(inst, "grpc.ExecutePlan") + attrs = rec["attrs"] + assert attrs["kind"] == "unary_stream" + assert attrs["layer"] == "outer" + assert attrs["operation_id"] == "op-1" + assert attrs["session_id"] == "sess-1" + assert attrs["request_bytes"] == execute_request().ByteSize() + assert attrs["response_count"] == 3 + assert attrs["response_bytes"] == sum(m.ByteSize() for m in msgs) + assert attrs["grpc_code"] == "OK" + assert [e["name"] for e in rec["events"]] == ["first_response", "status"] + assert rec["status"] == "ok" + + +def test_unary_stream_error_mid_stream(inst): + fake = FakeChannel() + error = FakeRpcError(grpc.StatusCode.INTERNAL) + fake.handlers[EXECUTE] = lambda req: FakeStreamCall(responses(1), error) + channel = wrap_channel(fake, inst) + call = channel.unary_stream(EXECUTE)(execute_request()) + next(call) + with pytest.raises(FakeRpcError) as info: + next(call) + assert info.value is error # the real call is what pyspark inspects + rec = one(inst, "grpc.ExecutePlan") + assert rec["status"] == "error" + assert rec["error"]["grpc_code"] == "INTERNAL" + assert rec["attrs"]["response_count"] == 1 + + +def test_unary_stream_cancel_ends_span(inst): + fake = FakeChannel() + fake.handlers[EXECUTE] = lambda req: FakeStreamCall(responses(5)) + channel = wrap_channel(fake, inst) + call = channel.unary_stream(EXECUTE)(execute_request()) + next(call) + assert call.cancel() is True + rec = one(inst, "grpc.ExecutePlan") + assert rec["attrs"]["cancelled"] is True + assert rec["attrs"]["response_count"] == 1 + + +def test_unary_stream_status_callback_from_grpc_thread(inst): + """The status event comes from add_callback, before the consumer + observes the end of the stream.""" + fake = FakeChannel() + stream = FakeStreamCall(responses(1)) + fake.handlers[EXECUTE] = lambda req: stream + channel = wrap_channel(fake, inst) + call = channel.unary_stream(EXECUTE)(execute_request()) + next(call) + stream._finish() # gRPC reports completion + assert spans(inst, "grpc.ExecutePlan") == [] + with pytest.raises(StopIteration): + next(call) + rec = one(inst, "grpc.ExecutePlan") + assert [e["name"] for e in rec["events"]] == ["first_response", "status"] + + +def test_stream_unary_counts_requests(inst): + fake = FakeChannel() + seen = [] + + def handler(request_iterator): + seen.extend(request_iterator) + return proto.AddArtifactsResponse() + + fake.handlers[ARTIFACTS] = handler + channel = wrap_channel(fake, inst) + chunks = [proto.AddArtifactsRequest(session_id="sess-1") for _ in range(4)] + channel.stream_unary(ARTIFACTS)(iter(chunks)) + assert seen == chunks + rec = one(inst, "grpc.AddArtifacts") + assert rec["attrs"]["kind"] == "stream_unary" + assert rec["attrs"]["request_count"] == 4 + assert rec["attrs"]["request_bytes"] == sum(c.ByteSize() for c in chunks) + assert rec["attrs"]["session_id"] == "sess-1" + assert rec["attrs"]["response_count"] == 1 + + +def test_stream_stream_span(inst): + fake = FakeChannel() + fake.handlers[EXECUTE] = lambda reqs: FakeStreamCall( + responses(len(list(reqs))) + ) + channel = wrap_channel(fake, inst) + call = channel.stream_stream(EXECUTE)(iter([execute_request()] * 2)) + assert len(list(call)) == 2 + rec = one(inst, "grpc.ExecutePlan") + assert rec["attrs"]["kind"] == "stream_stream" + assert rec["attrs"]["request_count"] == 2 + assert rec["attrs"]["response_count"] == 2 + + +def test_wrap_channel_disabled_is_passthrough(): + fake = FakeChannel() + assert wrap_channel(fake, None) is fake + assert wrap_channel(fake, NOOP_RECORDER, "inner") is fake + with pytest.raises(ValueError): + GrpcInterceptor(Recorder(lambda r: None), "middle") + + +def test_spans_nest_under_the_current_span(inst): + fake = FakeChannel() + channel = wrap_channel(fake, inst) + with inst.recorder.span("client.to_pandas"): + channel.unary_unary(CONFIG)(proto.ConfigRequest()) + assert ( + one(inst, "grpc.Config")["parent_id"] + == one(inst, "client.to_pandas")["span_id"] + ) + + +@pytest.fixture +def proxied(inst): + """A ProxiedChannel over a FakeChannel, with instrumentation.""" + fake = FakeChannel() + with ( + mock.patch.object(core.proxy, "ManagedSparkSessionProxy") as proxy_cls, + mock.patch.object(core, "DefaultChannelBuilder") as builder_cls, + ): + proxy_cls.return_value.port = 1 + builder_cls.return_value.toChannel.return_value = fake + yield fake, proxy_cls + + +def test_two_layers_with_guard(inst, proxied): + fake, proxy_cls = proxied + fake.handlers[EXECUTE] = lambda req: FakeStreamCall(responses(2)) + guard_calls = [] + + def guard(): + with inst.recorder.span( + "dataproc.get_session", reason="rpc_guard", new_client=True + ): + guard_calls.append(1) + return True + + builder = core.ManagedSparkChannelBuilder( + "sc://example.com:443/;use_ssl=true", + is_active_callback=guard, + instrumentation=inst, + ) + channel = builder.toChannel() + proxy_cls.assert_called_once_with(0, "example.com", recorder=inst.recorder) + list(channel.unary_stream(EXECUTE)(execute_request())) + + outer = [ + s + for s in spans(inst, "grpc.ExecutePlan") + if s["attrs"]["layer"] == "outer" + ] + inner = [ + s + for s in spans(inst, "grpc.ExecutePlan") + if s["attrs"]["layer"] == "inner" + ] + assert len(outer) == 1 and len(inner) == 1 + assert inner[0]["parent_id"] == outer[0]["span_id"] + assert outer[0]["parent_id"] is None + assert guard_calls == [1] + assert outer[0]["attrs"]["guard_ns"] > 0 + assert "guard_ns" not in inner[0]["attrs"] + guard_span = one(inst, "dataproc.get_session") + assert guard_span["parent_id"] == outer[0]["span_id"] + assert guard_span["attrs"]["reason"] == "rpc_guard" + assert outer[0]["attrs"]["response_count"] == 2 + assert inner[0]["attrs"]["response_count"] == 2 + assert outer[0]["duration_ns"] >= inner[0]["duration_ns"] + + +def test_guard_failure_is_an_error_on_the_outer_span(proxied): + fake, _ = proxied + snapshots = [] + inst = Instrumentation( + [ + InMemorySink(), + CallbackSink(lambda r: snapshots.append(json.loads(json.dumps(r)))), + ] + ) + channel = core.ManagedSparkChannelBuilder( + "sc://example.com:443", + is_active_callback=lambda: False, + instrumentation=inst, + ).toChannel() + with pytest.raises(RuntimeError, match="Session not active"): + channel.unary_unary(CONFIG)(proto.ConfigRequest()) + rec = one(inst, "grpc.Config") + assert rec["attrs"]["layer"] == "outer" + assert rec["status"] == "error" + assert rec["error"]["type"] == "RuntimeError" + assert rec["attrs"]["guard_ns"] >= 0 + assert spans(inst, "grpc.Config") == [rec] # inner never ran + (emitted,) = [r for r in snapshots if r["type"] == "span"] + assert emitted["attrs"]["guard_ns"] >= 0 + + +def test_ended_span_ignores_late_set_and_record_exception(inst): + span = inst.recorder.start_span("done", initial=1) + span.end() + span.set(initial=2, extra=3) + span.record_exception(ValueError("too late")) + rec = one(inst, "done") + assert rec["status"] == "ok" + assert rec["error"] is None + assert rec["attrs"] == {"run_id": "r1", "initial": 1} + + +def test_disabled_channel_builder_installs_nothing(proxied): + fake, proxy_cls = proxied + cb = mock.Mock(return_value=True) + channel = core.ManagedSparkChannelBuilder( + "sc://example.com:443", is_active_callback=cb + ).toChannel() + assert isinstance(channel, core.ProxiedChannel) + assert channel._is_active_callback is cb # not wrapped by timed_guard + assert channel._wrapped is fake # no inner interceptor + proxy_cls.assert_called_once_with(0, "example.com", recorder=NOOP_RECORDER) + + +def test_timed_guard(): + assert timed_guard(None) is None + guard = timed_guard(lambda: True) + assert guard() is True + assert guard.__name__ == "" + + +# -------------------------------------------------------------------------- +# Client wrappers +# -------------------------------------------------------------------------- + + +class FakeClient: + + def __init__(self): + self.table = mock.Mock(num_rows=7, num_columns=2) + + def _execute_and_fetch(self, req, observations, self_destruct=False): + return self.table, None, [], [], None + + def to_pandas(self, plan, observations): + table, *_ = self._execute_and_fetch(plan, observations) + return "pdf", "ei" + + +def test_client_to_pandas_wrappers(inst): + client = FakeClient() + install_client_wrappers(client, inst.recorder) + install_client_wrappers(client, inst.recorder) # idempotent + assert client.to_pandas("plan", {}) == ("pdf", "ei") + to_pandas = one(inst, "client.to_pandas") + conv = one(inst, "client.arrow_to_pandas") + assert conv["parent_id"] == to_pandas["span_id"] + assert conv["attrs"]["rows"] == 7 and conv["attrs"]["columns"] == 2 + assert conv["start_mono_ns"] >= to_pandas["start_mono_ns"] + assert ( + conv["start_mono_ns"] + conv["duration_ns"] + <= to_pandas["start_mono_ns"] + to_pandas["duration_ns"] + ) + # A fetch outside to_pandas does not start a conversion span. + client._execute_and_fetch("plan", {}) + assert len(spans(inst, "client.arrow_to_pandas")) == 1 + + uninstall_client_wrappers(client) + uninstall_client_wrappers(client) + assert "to_pandas" not in client.__dict__ + assert "_execute_and_fetch" not in client.__dict__ + client.to_pandas("plan", {}) + assert len(spans(inst, "client.to_pandas")) == 1 + + +def test_client_to_pandas_conversion_error_marks_arrow_to_pandas_error(inst): + class FailingConversionClient(FakeClient): + + def to_pandas(self, plan, observations): + self._execute_and_fetch(plan, observations) + raise ValueError("bad arrow conversion") + + client = FailingConversionClient() + install_client_wrappers(client, inst.recorder) + try: + with pytest.raises(ValueError, match="bad arrow conversion"): + client.to_pandas("plan", {}) + finally: + uninstall_client_wrappers(client) + + to_pandas = one(inst, "client.to_pandas") + conv = one(inst, "client.arrow_to_pandas") + assert to_pandas["status"] == "error" + assert to_pandas["error"]["type"] == "ValueError" + assert conv["status"] == "error" + assert conv["error"]["type"] == "ValueError" + assert conv["parent_id"] == to_pandas["span_id"] + + +def test_client_class_hooks(inst): + import pyarrow as pa + from pyspark.sql.connect.plan import Range + from pyspark.sql.conversion import ArrowTableToRowsConversion + from pyspark.sql.types import LongType, StructField, StructType + + client = FakeClient() + other = FakeClient() + install_client_wrappers(client, inst.recorder) + try: + plan = Range(start=0, end=10, step=1, num_partitions=1).to_proto(client) + assert plan.root.HasField("range") + rec = one(inst, "client.plan_to_proto") + assert rec["attrs"]["plan_bytes"] == plan.ByteSize() + + # Another client (no recorder) goes through the hook untraced. + Range(start=0, end=1, step=1, num_partitions=1).to_proto(other) + assert len(spans(inst, "client.plan_to_proto")) == 1 + + table = pa.table({"a": [1, 2, 3]}) + schema = StructType([StructField("a", LongType())]) + rows = ArrowTableToRowsConversion.convert(table, schema) + assert [r.a for r in rows] == [1, 2, 3] + assert one(inst, "client.arrow_to_rows")["attrs"]["rows"] == 3 + finally: + uninstall_client_wrappers(client) + + # With no instrumented client registered the hooks record nothing. + ArrowTableToRowsConversion.convert(table, schema) + assert len(spans(inst, "client.arrow_to_rows")) == 1 + assert _client._active_recorder() is NOOP_RECORDER + + +def test_install_is_a_noop_when_disabled(): + client = FakeClient() + install_client_wrappers(client, NOOP_RECORDER) + assert "to_pandas" not in client.__dict__ + assert not hasattr(client, _client._RECORDER_ATTR) + + +# -------------------------------------------------------------------------- +# Session wiring (mocked Dataproc API, no network) +# -------------------------------------------------------------------------- + + +def _session_response(session_id, uuid): + from google.cloud.dataproc_v1 import Session + + response = Session() + response.name = ( + f"projects/test-project/locations/test-region/sessions/{session_id}" + ) + response.uuid = uuid + response.state = Session.State.ACTIVE + response.runtime_config.version = "3.0" + response.session_template = "projects/p/locations/r/sessionTemplates/t" + response.runtime_info.endpoints = { + "Spark Connect Server": "sc://spark-connect-server.example.com:443" + } + response.create_time = "2024-01-01T00:00:00Z" + response.state_time = "2024-01-01T00:01:00Z" + response.state_history = [ + { + "state": Session.State.CREATING, + "state_start_time": "2024-01-01T00:00:00Z", + } + ] + return response + + +@pytest.fixture +def dataproc_mocks(monkeypatch): + """Everything getOrCreate()/stop() touch, without a network.""" + from google.cloud.dataproc_v1 import Session + from pyspark.sql.connect.client.core import ConfigResult + from pyspark.sql.connect.proto import ConfigResponse + + from google.cloud.managed_spark_connect import ManagedSparkSession + + for key in list(os.environ): + if key.startswith("MANAGED_SPARK_CONNECT_"): + monkeypatch.delenv(key) + session_id = "sc-20240702-103952-abcdef" + uuid = "c002e4ef-fe5e-41a8-a157-160aa73e4f7f" + response = _session_response(session_id, uuid) + + with ( + mock.patch("google.auth.default") as auth_default, + mock.patch( + "google.cloud.dataproc_v1.SessionControllerClient" + ) as client_cls, + mock.patch( + "pyspark.sql.connect.client.SparkConnectClient.config", + return_value=ConfigResult.fromProto(ConfigResponse()), + ), + mock.patch.object( + ManagedSparkSession.Builder, + "generate_session_id", + return_value=session_id, + ), + mock.patch( + "google.cloud.managed_spark_connect.session.is_s8s_session_active", + return_value=True, + ), + mock.patch( + "google.cloud.managed_spark_connect.session._server_spark_version", + return_value=None, + ), + mock.patch( + "google.cloud.managed_spark_connect.environment.get_client_environment_label", + return_value="unknown", + ), + mock.patch("time.sleep"), + ): + cred = mock.MagicMock() + cred.token = "token" + auth_default.return_value = (cred, "") + client = client_cls.return_value + operation = mock.Mock() + operation.result.return_value = response + client.create_session.return_value = operation + terminating = Session() + terminating.state = Session.State.TERMINATING + client.get_session.return_value = terminating + yield { + "client": client, + "session_id": session_id, + "uuid": uuid, + "response": response, + } + + +def _by_name(inst): + out = {} + for rec in spans(inst): + out.setdefault(rec["name"], []).append(rec) + return out + + +def test_session_builder_opt_in(dataproc_mocks): + from google.cloud.managed_spark_connect import ManagedSparkSession + + inst = Instrumentation([InMemorySink()], {"run_id": "r1"}) + session = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .instrumentation(inst) + .getOrCreate() + ) + try: + assert session.instrumentation is inst + assert getattr(session.client, _client._RECORDER_ATTR) is inst.recorder + assert inst.recorder.trace_id == dataproc_mocks["uuid"] + + by_name = _by_name(inst) + root = by_name["session.get_or_create"][0] + assert root["parent_id"] is None + assert root["attrs"] == { + "run_id": "r1", + "mode": "create", + "session_id": dataproc_mocks["session_id"], + "uuid": dataproc_mocks["uuid"], + "runtime_version": "3.0", + "template": "projects/p/locations/r/sessionTemplates/t", + } + for child in ( + "session.create_rpc", + "session.wait_active", + "session.connect", + ): + (rec,) = by_name[child] + assert rec["parent_id"] == root["span_id"], child + wait = by_name["session.wait_active"][0]["attrs"] + assert wait["poll_interval_s"] == 5.0 + assert wait["create_time"].startswith("2024-01-01T00:00:00") + assert wait["state"] == "ACTIVE" + assert json.loads(wait["state_history"])[0]["state"] == "CREATING" + (start,) = by_name["proxy.start"] + assert start["parent_id"] == by_name["session.connect"][0]["span_id"] + assert start["attrs"]["port"] > 0 + assert "session.reuse" not in by_name + # Everything after connect carries the session uuid as trace id. + assert start["trace_id"] == dataproc_mocks["uuid"] + assert root["trace_id"] == dataproc_mocks["uuid"] + finally: + session.stop() + + by_name = _by_name(inst) + (stop,) = by_name["session.stop"] + assert stop["attrs"]["terminate"] is True + assert stop["attrs"]["session_id"] == dataproc_mocks["session_id"] + assert stop["trace_id"] == dataproc_mocks["uuid"] + assert by_name["session.terminate_rpc"][0]["parent_id"] == stop["span_id"] + (wait,) = by_name["session.wait_terminating"] + assert wait["parent_id"] == stop["span_id"] + assert wait["attrs"]["state"] == "TERMINATING" + (poll,) = by_name["dataproc.get_session"] + assert poll["parent_id"] == wait["span_id"] + assert poll["attrs"] == { + "run_id": "r1", + "reason": "stop", + "new_client": False, + } + dataproc_mocks["client"].get_session.assert_called_once() + # stop() closes the instrumentation and detaches the client wrappers. + assert inst.recorder.trace_id == "" + assert not hasattr(session.client, _client._RECORDER_ATTR) + + +def test_session_env_opt_in(dataproc_mocks, monkeypatch, tmp_path): + from google.cloud.managed_spark_connect import ManagedSparkSession + + path = tmp_path / "spans.jsonl" + monkeypatch.setenv(inst_lib.ENV_FILE, str(path)) + session = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + try: + assert session.instrumentation is not None + assert [type(s) for s in session.instrumentation.sinks] == [ + JsonlFileSink + ] + finally: + session.stop() + records = [json.loads(l) for l in path.read_text().splitlines()] + assert records[0]["type"] == "meta" + names = {r["name"] for r in records if r["type"] == "span"} + assert {"session.get_or_create", "session.connect", "session.stop"} <= names + + +def test_session_disabled_by_default(dataproc_mocks): + from google.cloud.managed_spark_connect import ManagedSparkSession + + session = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + try: + assert session.instrumentation is None + assert not hasattr(session.client, _client._RECORDER_ATTR) + assert isinstance(session.client._channel, core.ProxiedChannel) + assert "to_pandas" not in session.client.__dict__ + finally: + session.stop() + + +def test_session_reuse_spans(dataproc_mocks): + from google.cloud.managed_spark_connect import ManagedSparkSession + + inst = Instrumentation([InMemorySink()]) + builder = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .instrumentation(inst) + ) + first = builder.getOrCreate() + try: + with mock.patch( + "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", + return_value=dataproc_mocks["response"], + ): + second = builder.getOrCreate() + assert second is first + by_name = _by_name(inst) + roots = by_name["session.get_or_create"] + assert [r["attrs"]["mode"] for r in roots] == ["create", "reuse"] + (reuse,) = by_name["session.reuse"] + assert reuse["parent_id"] == roots[1]["span_id"] + assert reuse["attrs"]["found"] is True + (lookup,) = by_name["dataproc.get_session"] + assert lookup["parent_id"] == reuse["span_id"] + assert lookup["attrs"] == {"reason": "reuse", "new_client": True} + finally: + first.stop() + + +def test_session_reuse_fresh_builder_reuses_instrumentation_and_sets_trace_id( + dataproc_mocks, monkeypatch, tmp_path +): + from google.cloud.managed_spark_connect import ManagedSparkSession + + # 1. Env-configured instrumentation: fresh builder reuses active session's + # Instrumentation without emitting a duplicate meta record or orphaning + # the reuse span. + path = tmp_path / "reuse.jsonl" + monkeypatch.setenv(inst_lib.ENV_FILE, str(path)) + first = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + try: + with mock.patch( + "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", + return_value=dataproc_mocks["response"], + ): + second = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + assert second is first + finally: + first.stop() + + records = [json.loads(line) for line in path.read_text().splitlines()] + assert [r["type"] for r in records].count("meta") == 1 + get_or_create = [ + r + for r in records + if r["type"] == "span" and r["name"] == "session.get_or_create" + ] + assert [r["attrs"]["mode"] for r in get_or_create] == ["create", "reuse"] + assert all(r["trace_id"] == dataproc_mocks["uuid"] for r in get_or_create) + monkeypatch.delenv(inst_lib.ENV_FILE) + + # 2. Programmatic instrumentation: fresh builder without .instrumentation() + # records reuse on the active session's Instrumentation; explicit new + # Instrumentation on reuse still gets trace_id populated. + inst1 = Instrumentation([InMemorySink()]) + inst2 = Instrumentation([InMemorySink()]) + session = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .instrumentation(inst1) + .getOrCreate() + ) + try: + with mock.patch( + "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", + return_value=dataproc_mocks["response"], + ): + reused_default = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + reused_explicit = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .instrumentation(inst2) + .getOrCreate() + ) + assert reused_default is session + assert reused_explicit is session + assert [ + r["attrs"]["mode"] for r in _by_name(inst1)["session.get_or_create"] + ] == ["create", "reuse"] + assert inst2.recorder.trace_id == dataproc_mocks["uuid"] + for name in ( + "session.get_or_create", + "session.reuse", + "dataproc.get_session", + ): + (rec,) = _by_name(inst2)[name] + assert rec["trace_id"] == dataproc_mocks["uuid"], name + finally: + session.stop() + + +def test_session_attach_when_local_session_missing_connects_before_versions( + dataproc_mocks, +): + from google.cloud.managed_spark_connect import ManagedSparkSession + + seen_sessions = [] + orig_versions = ManagedSparkSession.Builder._session_versions + + def spy_versions(sess): + seen_sessions.append(sess) + return orig_versions(sess) + + ManagedSparkSession._active_s8s_session_id = dataproc_mocks["session_id"] + session = None + try: + with ( + mock.patch( + "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", + return_value=dataproc_mocks["response"], + ), + mock.patch.object( + ManagedSparkSession.Builder, + "_wait_for_session_available", + return_value=dataproc_mocks["response"], + ), + mock.patch.object( + ManagedSparkSession.Builder, + "_session_versions", + staticmethod(spy_versions), + ), + ): + session = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + assert seen_sessions == [session] + assert session is not None + finally: + if session is not None: + session.stop() + ManagedSparkSession._active_s8s_session_id = None + + +# -------------------------------------------------------------------------- +# Real gRPC round trip (in-process server on 127.0.0.1, no network) +# -------------------------------------------------------------------------- + + +@pytest.fixture +def grpc_server(): + from concurrent import futures + + from pyspark.sql.connect.proto import base_pb2_grpc + + class Service(base_pb2_grpc.SparkConnectServiceServicer): + + def ExecutePlan(self, request, context): + yield proto.ExecutePlanResponse(session_id=request.session_id) + if request.operation_id == "fail": + context.abort(grpc.StatusCode.INTERNAL, "boom") + for _ in range(2): + yield proto.ExecutePlanResponse( + session_id=request.session_id, + operation_id=request.operation_id, + ) + + def Config(self, request, context): + return proto.ConfigResponse(session_id=request.session_id) + + server = grpc.server(futures.ThreadPoolExecutor(max_workers=4)) + base_pb2_grpc.add_SparkConnectServiceServicer_to_server(Service(), server) + port = server.add_insecure_port("127.0.0.1:0") + server.start() + try: + yield port + finally: + server.stop(None) + + +def _local_channel(port): + return grpc.insecure_channel( + f"127.0.0.1:{port}", options=[("grpc.enable_http_proxy", 0)] + ) + + +def test_real_grpc_round_trip_through_proxied_channel(inst, grpc_server): + from pyspark.sql.connect.proto import base_pb2_grpc + + guard = mock.Mock(return_value=True) + with ( + mock.patch.object(core.proxy, "ManagedSparkSessionProxy") as proxy_cls, + mock.patch.object(core, "DefaultChannelBuilder") as builder_cls, + ): + proxy_cls.return_value.port = 1 + builder_cls.return_value.toChannel.return_value = _local_channel( + grpc_server + ) + channel = core.ManagedSparkChannelBuilder( + "sc://example.com:443", + is_active_callback=guard, + instrumentation=inst, + ).toChannel() + stub = base_pb2_grpc.SparkConnectServiceStub(channel) + try: + assert ( + stub.Config(proto.ConfigRequest(session_id="s1")).session_id == "s1" + ) + call = stub.ExecutePlan( + proto.ExecutePlanRequest(session_id="s1", operation_id="op1") + ) + assert len(list(call)) == 3 + assert call.code() == grpc.StatusCode.OK + with pytest.raises(grpc.RpcError) as info: + list( + stub.ExecutePlan( + proto.ExecutePlanRequest( + session_id="s1", operation_id="fail" + ) + ) + ) + assert info.value.code() == grpc.StatusCode.INTERNAL + finally: + channel.close() + + assert guard.call_count == 3 + by_name = _by_name(inst) + config = by_name["grpc.Config"] + assert [s["attrs"]["layer"] for s in config] == ["inner", "outer"] + assert config[0]["parent_id"] == config[1]["span_id"] + assert config[1]["attrs"]["grpc_code"] == "OK" + assert config[1]["attrs"]["guard_ns"] > 0 + + execute = by_name["grpc.ExecutePlan"] + ok = [s for s in execute if s["attrs"]["operation_id"] == "op1"] + failed = [s for s in execute if s["attrs"]["operation_id"] == "fail"] + assert len(ok) == 2 and len(failed) == 2 + for rec in ok: + assert rec["attrs"]["response_count"] == 3 + assert rec["attrs"]["grpc_code"] == "OK" + assert [e["name"] for e in rec["events"]] == [ + "first_response", + "status", + ] + for rec in failed: + assert rec["status"] == "error" + assert rec["error"]["grpc_code"] == "INTERNAL" + assert rec["error"]["message"] == "boom" + assert rec["attrs"]["response_count"] == 1 + + +# -------------------------------------------------------------------------- +# Abandoned streams, guard idempotence, attr snapshots, status end time +# -------------------------------------------------------------------------- + + +def test_abandoned_stream_span_is_emitted_when_iterator_is_dropped(inst): + """pyspark stops reading an ExecutePlan stream at result_complete and + drops the iterator without draining it or calling cancel().""" + fake = FakeChannel() + stream = FakeStreamCall(responses(3)) + fake.handlers[EXECUTE] = lambda req: stream + channel = wrap_channel(wrap_channel(fake, inst, "inner"), inst, "outer") + call = channel.unary_stream(EXECUTE)(execute_request()) + next(call) + next(call) + stream._finish() # the server is done; status arrives on the grpc thread + assert spans(inst, "grpc.ExecutePlan") == [] + del call # consumer walks away + recs = spans(inst, "grpc.ExecutePlan") + assert [r["attrs"]["layer"] for r in recs] == ["inner", "outer"] + for rec in recs: + assert rec["attrs"]["response_count"] == 2 + assert rec["attrs"]["grpc_code"] == "OK" + assert [e["name"] for e in rec["events"]] == [ + "first_response", + "status", + ] + assert recs[0]["parent_id"] == recs[1]["span_id"] + assert inst.recorder.inflight() == [] + assert inst.wait_idle(0) is True + + +def test_abandoned_active_stream_does_not_block(inst): + fake = FakeChannel() + fake.handlers[EXECUTE] = lambda req: FakeStreamCall(responses(3)) + channel = wrap_channel(fake, inst) + call = channel.unary_stream(EXECUTE)(execute_request()) + next(call) + del call # RPC still active: code() must not be consulted (it blocks) + (rec,) = spans(inst, "grpc.ExecutePlan") + assert rec["attrs"]["response_count"] == 1 + assert "grpc_code" not in rec["attrs"] + assert [e["name"] for e in rec["events"]] == ["first_response"] + + +def test_late_status_callback_does_not_touch_an_ended_span(inst): + fake = FakeChannel() + stream = FakeStreamCall(responses(1)) + fake.handlers[EXECUTE] = lambda req: stream + call = wrap_channel(fake, inst).unary_stream(EXECUTE)(execute_request()) + next(call) + del call + (rec,) = spans(inst, "grpc.ExecutePlan") + before = json.dumps(rec, sort_keys=True) + stream._finish() # callback fires after the span ended + assert json.dumps(rec, sort_keys=True) == before + + +def test_timed_guard_is_idempotent(): + inner = timed_guard(lambda: True) + assert timed_guard(inner) is inner + assert timed_guard(timed_guard(inner)) is inner + + +def test_guard_wrapped_twice_counts_once(inst, proxied, monkeypatch): + fake, _ = proxied + clock = [1_000_000_000] + + def slow_guard(): + clock[0] += 20_000_000 + return True + + channel = core.ManagedSparkChannelBuilder( + "sc://example.com:443", + is_active_callback=timed_guard(slow_guard), # as scbench does + instrumentation=inst, + ).toChannel() + monkeypatch.setattr("time.perf_counter_ns", lambda: clock[0]) + channel.unary_unary(CONFIG)(proto.ConfigRequest()) + rec = [ + s for s in spans(inst, "grpc.Config") if s["attrs"]["layer"] == "outer" + ][0] + assert rec["attrs"]["guard_ns"] == 20_000_000 + + +def test_global_attributes_are_snapshotted_when_a_span_starts(inst): + early = inst.recorder.start_span("early") + inst.set_attribute("trial_id", "t1") + during = inst.recorder.start_span("during") + inst.remove_attribute("trial_id") + inst.set_attribute("trial_id", "t2") + early.end() + during.end() + assert "trial_id" not in one(inst, "early")["attrs"] + assert one(inst, "during")["attrs"]["trial_id"] == "t1" + # Events are stamped when emitted. + inst.recorder.event("e") + assert [r for r in inst.records if r["type"] == "event"][0]["attrs"][ + "trial_id" + ] == "t2" + + +def test_stream_span_ends_at_status_not_at_release(inst): + """pyspark keeps the finished iterator alive until its ReleaseExecute + has been queued; the span must still end when the status arrived.""" + import time + + fake = FakeChannel() + stream = FakeStreamCall(responses(3)) + fake.handlers[EXECUTE] = lambda req: stream + call = wrap_channel(fake, inst).unary_stream(EXECUTE)(execute_request()) + next(call) + next(call) + stream._finish() # status from the grpc thread: the RPC is over + time.sleep(0.3) # consumer still holds the iterator + released_mono_ns = time.perf_counter_ns() + del call + (rec,) = spans(inst, "grpc.ExecutePlan") + status = [e for e in rec["events"] if e["name"] == "status"][0] + assert ( + rec["duration_ns"] + <= released_mono_ns - rec["start_mono_ns"] - 200_000_000 + ) + assert abs(rec["duration_ns"] - status["offset_ns"]) < 50_000_000 + assert rec["attrs"]["response_count"] == 2 + + +def test_span_end_at_explicit_time(inst): + import time + + span = inst.recorder.start_span("x") + then = time.perf_counter_ns() + time.sleep(0.05) + span.end(end_mono_ns=then) + assert one(inst, "x")["duration_ns"] == then - span.start_mono_ns + early = inst.recorder.start_span("y") + early.end(end_mono_ns=early.start_mono_ns - 10) + assert one(inst, "y")["duration_ns"] == 0 + + +# -------------------------------------------------------------------------- +# In-flight tracking +# -------------------------------------------------------------------------- + + +def test_inflight_and_wait_idle(inst): + import time + + assert inst.recorder.inflight() == [] + assert inst.wait_idle(0) is True + span = inst.recorder.start_span("grpc.ReleaseExecute") + other = inst.recorder.start_span("proxy.connection") + assert [s.name for s in inst.recorder.inflight("grpc.")] == [ + "grpc.ReleaseExecute" + ] + assert len(inst.recorder.inflight()) == 2 + assert inst.wait_idle(0.05) is False + assert inst.wait_idle(0.05, prefix="client.") is True + + def finish_later(): + time.sleep(0.05) + span.end() + + t = threading.Thread(target=finish_later) + t.start() + assert inst.wait_idle(5) is True # woken by span.end(), not the timeout + t.join() + assert inst.recorder.inflight("grpc.") == [] + other.end() + assert inst.recorder.inflight() == [] + assert NOOP_RECORDER.inflight() == [] + assert NOOP_RECORDER.wait_idle(0) is True diff --git a/tests/unit/test_proxy.py b/tests/unit/test_proxy.py index 23339391..8aa8172b 100644 --- a/tests/unit/test_proxy.py +++ b/tests/unit/test_proxy.py @@ -149,3 +149,191 @@ def test_proxy_with_timeouts(client_wait, proxy_server_conn, test_message): retry_on_timeouts(proxy_server_conn.recv, 1024).decode() ) assert "\n".join(sent) == "\n".join(received) + + +# -------------------------------------------------------------------------- +# Instrumentation of the proxy (no network) +# -------------------------------------------------------------------------- + + +def _spans(inst, name): + return [ + r for r in inst.records if r["type"] == "span" and r["name"] == name + ] + + +def test_connect_sockets_counts_traffic(): + from google.cloud.managed_spark_connect.instrumentation import ( + InMemorySink, + Instrumentation, + ) + + inst = Instrumentation([InMemorySink()]) + client, proxy_side = socket.socketpair() + bridge_side, backend = socket.socketpair() + span = inst.recorder.start_span("proxy.connection", conn=1) + t = threading.Thread( + target=connect_sockets, + args=[1, proxy_side, bridge_side, span], + daemon=True, + ) + t.start() + client.sendall(b"hello") + assert backend.recv(1024) == b"hello" + backend.sendall(b"world!!") + assert client.recv(1024) == b"world!!" + client.close() + backend.close() + t.join(timeout=5) + assert not t.is_alive() + span.end() + (rec,) = _spans(inst, "proxy.connection") + assert rec["attrs"] == { + "conn": 1, + "bytes_out": 5, + "frames_out": 1, + "bytes_in": 7, + "frames_in": 1, + } + + +def test_connect_tcp_bridge_spans(): + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + from google.cloud.managed_spark_connect.instrumentation import ( + InMemorySink, + Instrumentation, + ) + + inst = Instrumentation([InMemorySink()]) + creds = mock.Mock(token="tok") + with ( + mock.patch.object( + proxy.googleauth, "default", return_value=(creds, None) + ), + mock.patch.object(proxy.websocketclient, "connect") as connect, + ): + ws = proxy.connect_tcp_bridge("backend.example.com", inst.recorder, 3) + assert ws is connect.return_value + creds.refresh.assert_called_once() + (parent,) = _spans(inst, "proxy.ws_connect") + assert parent["attrs"] == {"conn": 3, "host": "backend.example.com"} + for name in ("auth.default", "auth.refresh", "ws.handshake"): + (child,) = _spans(inst, name) + assert child["parent_id"] == parent["span_id"] + + +def test_forward_connection_ends_span_on_failure(): + import time + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + from google.cloud.managed_spark_connect.instrumentation import ( + InMemorySink, + Instrumentation, + ) + + inst = Instrumentation([InMemorySink()]) + conn = mock.MagicMock() + accepted = (time.perf_counter_ns() - 5_000_000, time.time_ns()) + with mock.patch.object( + proxy, "connect_tcp_bridge", side_effect=OSError("no route") + ): + with pytest.raises(OSError): + proxy.forward_connection( + 1, conn, ("127.0.0.1", 1), "h", inst.recorder, accepted + ) + (rec,) = _spans(inst, "proxy.connection") + assert rec["status"] == "error" + assert rec["error"]["type"] == "OSError" + assert rec["attrs"]["conn"] == 1 + assert rec["start_mono_ns"] == accepted[0] + assert rec["duration_ns"] >= 5_000_000 + + +def test_forward_connection_nests_ws_connect_and_counts(): + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + from google.cloud.managed_spark_connect.instrumentation import ( + InMemorySink, + Instrumentation, + ) + + inst = Instrumentation([InMemorySink()]) + client, proxy_side = socket.socketpair() + bridge_side, backend = socket.socketpair() + creds = mock.Mock(token="tok") + + class FakeWs: + + def __enter__(self): + return self + + def __exit__(self, *args): + bridge_side.close() + return False + + def recv(self, timeout=None): + data = bridge_side.recv(1024) + if not data: + raise ConnectionError("closed") + return data.hex() + + def send(self, msg): + bridge_side.sendall(bytes.fromhex(msg)) + + def close(self): + bridge_side.close() + + with ( + mock.patch.object( + proxy.googleauth, "default", return_value=(creds, None) + ), + mock.patch.object( + proxy.websocketclient, "connect", return_value=FakeWs() + ), + ): + t = threading.Thread( + target=proxy.forward_connection, + args=[2, proxy_side, ("127.0.0.1", 1), "h", inst.recorder], + daemon=True, + ) + t.start() + client.sendall(b"ping") + assert backend.recv(1024) == b"ping" + backend.sendall(b"pong!") + assert client.recv(1024) == b"pong!" + client.close() + backend.close() + t.join(timeout=5) + assert not t.is_alive() + + (conn_span,) = _spans(inst, "proxy.connection") + (ws,) = _spans(inst, "proxy.ws_connect") + assert ws["parent_id"] == conn_span["span_id"] + assert _spans(inst, "auth.default")[0]["parent_id"] == ws["span_id"] + assert conn_span["thread"] == ws["thread"] == t.name + assert conn_span["attrs"] == { + "conn": 2, + "bytes_out": 4, + "frames_out": 1, + "bytes_in": 5, + "frames_in": 1, + } + + +def test_proxy_start_span(): + from google.cloud.managed_spark_connect.client import proxy + from google.cloud.managed_spark_connect.instrumentation import ( + InMemorySink, + Instrumentation, + ) + + inst = Instrumentation([InMemorySink()]) + p = proxy.ManagedSparkSessionProxy(0, "h", recorder=inst.recorder) + p.start() + p.stop() + (rec,) = _spans(inst, "proxy.start") + assert rec["attrs"]["port"] == p.port > 0