feat(instrumentation): opt-in instrumentation of sessions, RPCs, proxy and client

.gitignore

@@ -9,3 +9,4 @@ __pycache__
 build/
 dist/
 
+.venv/

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

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":"<Dataproc session uuid, '' before it is known>",
+ "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.<Method>` (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.

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)