fix(session): notice, explain and replace a new session that fails to start

README.md

@@ -138,7 +138,8 @@ These environment variables change how the client behaves.
 |----------|---------|--------|
 | `MANAGED_SPARK_CONNECT_REPLAY_ON_POSITION_LOST` | `1` | When a read query (`collect`, `toPandas`, `count`, ...) fails with `INVALID_CURSOR.POSITION_NOT_AVAILABLE` because the server no longer buffers the result the client has to resume from, re-run it once under a new operation id and log a `WARNING`. Commands (writes, DDL) are never re-run, and neither are queries whose plan contains a UDF, UDTF, `mapInPandas`, `mapInArrow` or `applyInPandas`, or that come from `foreach` or `foreachPartition`, since re-running them re-runs any side effects of that code. UDFs only called by name (registered with `spark.udf.register` and called from SQL) aren't detected. Set to `0` to disable, e.g. for queries that are expensive or not deterministic, or that call side-effecting UDFs by name. Read when the session is created. |
 | `MANAGED_SPARK_CONNECT_WS_SILENCE_LIMIT_S` | `300` | Cut the proxy WebSocket connection if a query's result stream (a reattachable `ExecutePlan` or `ReattachExecute`) has been waiting this many seconds for its next response and nothing at all was received, so a wedged stream fails and pyspark reattaches instead of hanging. By default the server ends these streams every 2 minutes even when it has nothing to send, so that much silence means the connection is broken. Other RPCs, such as a long `AnalyzePlan`, and result streams that nobody is reading from don't count. The check runs on each 60 s WebSocket receive poll, so the cut happens between this limit and the limit plus 60 s. Set to `0` to disable. |
-| `MANAGED_SPARK_CONNECT_LIVENESS_TTL_SECONDS` | unset | Before an RPC, the client checks that the Dataproc session is still alive. Once a check finds it `ACTIVE`, that answer is reused for this many seconds: unset or negative reuses it until an RPC fails, `0` checks before every RPC. A failed RPC always triggers a fresh check. The session only counts as gone when Dataproc reports it not found or in a terminal state; when the check itself fails (for example during a control-plane outage), the RPC proceeds. |
+| `MANAGED_SPARK_CONNECT_LIVENESS_TTL_SECONDS` | unset | Before an RPC, the client checks that the Dataproc session is still alive. Once a check finds it `ACTIVE`, that answer is reused for this many seconds: unset or negative reuses it until an RPC fails, `0` checks before every RPC. A failed RPC always triggers a fresh check. The session only counts as gone when Dataproc reports it not found or in a terminal state; when the check itself fails (for example during a control-plane outage), the RPC proceeds. An RPC on a gone session raises `SessionNotActiveError` (a `RuntimeError`, from `google.cloud.managed_spark_connect.exceptions`) saying what Dataproc reported, e.g. `FAILED` and why. |
+| `MANAGED_SPARK_CONNECT_SESSION_START_RETRIES` | `1` | `getOrCreate()` makes a new session's first RPC itself, and if that fails and Dataproc reports the session gone (failed or terminated before it could be used), creates another session in its place, this many times, logging a `WARNING`. Only for generated session IDs: a session created with `sessionId()` raises `SessionNotActiveError` instead. `0` disables retrying. |
 | `GRPC_VERBOSITY` | `ERROR` | gRPC's own variable. If it is unset when this package is imported, and gRPC hasn't been imported yet, the package sets it to `ERROR` for the whole process to silence gRPC's informational fork messages. Set it yourself to override. |
 | `MANAGED_SPARK_CONNECT_LOG_LEVEL` | unset | A level name (`DEBUG`, `INFO`, ...) or a positive number for the `google.cloud.managed_spark_connect` loggers. If the root logger has no handlers when the package is imported, the library logs to stderr itself, and its records then don't reach root handlers configured later: configure logging before importing the package to send them elsewhere. The library never configures the root logger. Setting it also lets pyspark print its query-context JSON for failed queries (see `SPARK_CONNECT_LOG_LEVEL`). |
 | `SPARK_CONNECT_LOG_LEVEL` | unset | pyspark's own variable. Unless it or `MANAGED_SPARK_CONNECT_LOG_LEVEL` is set, creating a session silences pyspark's `SQLQueryContextLogger` and `DataFrameQueryContextLogger`, which otherwise print a JSON dump to stderr for every failed query. This applies to the whole process. |

docs/instrumentation.md

@@ -111,9 +111,9 @@ meta record carry them as of when they are written.
 
 | 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`, and `checked=false` when `GetSession` failed without a verdict and the session already in the process is reused unchecked |
+| `session.get_or_create` → `session.create_rpc`, `session.wait_active`, `session.wait_endpoint`, `session.connect`, `session.check_new` (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`, and `checked=false` when `GetSession` failed without a verdict and the session already in the process is reused unchecked; `session.check_new` (the new session's first RPC, which getOrCreate() makes itself) has `first_rpc_ok`, `timed_out` when the RPC was still running after the check timeout, and `dead` when that RPC failed and Dataproc was asked whether the session is gone |
 | `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) |
+| `dataproc.get_session` | `reason=rpc_guard\|reuse\|poll\|check_new\|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`, `plan_kind=root\|command` and `command_kind` (the command's oneof name, e.g. `sql_command`, `create_dataframe_view`) on `ExecutePlan`, `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` |
@@ -145,6 +145,7 @@ sessions; a session without an `Instrumentation` pays nothing for them.
 | `progress.bar_close` | The bar was closed and unregistered: by the handler on the final (`done=True`) message, or, if that never came, when the query's `_execute_and_fetch` returned or raised | `operation_id`, `completed_tasks`, `total_tasks`, `reason=done\|query_end\|replay` (`replay`: the query hit `INVALID_CURSOR.POSITION_NOT_AVAILABLE` and is re-run under a new operation id) |
 | `guard.result` | The session liveness guard ran (before an RPC, or after a failed one) and either had to ask the control plane or found the session inactive. A cached "active" verdict, the normal case for every RPC, emits nothing | `mode=pre_rpc\|on_error` (`direct` when called outside the channel), `cached` (answered from the liveness cache, no `GetSession` RPC), `active` (the verdict; `null` when the check raised), `error_type` (only when the check raised) |
 | `grpc.retry` | pyspark's retry loop is about to retry an RPC, after the policy's backoff. The failed RPC is the `grpc.*` span that ended just before, on the same `thread` | `attempt` (1 = first retry within that loop), `reason=error\|reattach` (`reattach`: the server ended the stream and pyspark reattaches), `grpc_code`, `error_type` |
+| `session.start_failed` | A session `getOrCreate()` just created never became usable: Dataproc reported it failed or terminated, either in the `CreateSession` result, while waiting for its endpoint, or after its first RPC failed (see `session.check_new`) | `detail` (what Dataproc reported, e.g. `Managed Spark Session sc-123 is FAILED: <reason>`), `attempt` (1 for the first session), `retry` (whether another session is created in its place: only for generated session IDs, up to `MANAGED_SPARK_CONNECT_SESSION_START_RETRIES` times) |
 | `reattach.position_lost` | A query's `_execute_and_fetch` (`collect`, `toPandas`, `count`, ...) failed with `INVALID_CURSOR.POSITION_NOT_AVAILABLE`: the server no longer buffers the response the client has to reattach from. A replay also logs a `WARNING` | `operation_id`, `replay` (whether the query is re-run once under a new operation id: only root (read) plans without inline UDFs or UDTFs, and not with `MANAGED_SPARK_CONNECT_REPLAY_ON_POSITION_LOST=0`; if the re-run fails the same way, a second event has its `operation_id` and `replay=false`), `new_operation_id` (only when `replay=true`) |
 
 Notes: `progress.*` events require the progress handler to be active,

google/cloud/managed_spark_connect/client/core.py

@@ -24,6 +24,7 @@
 from . import proxy
 from .proxy import get_default_credentials
 from .. import instrumentation as instrumentation_lib
+from ..exceptions import SessionNotActiveError
 from ..instrumentation._grpc import _StreamingCall
 
 logger = logging.getLogger(__name__)
@@ -468,9 +469,21 @@ def _check_active(self):
         self._guard_errors.success()
         if not active:
             logger.debug("Session is no longer active")
-            raise RuntimeError(
-                "Session not active. Please create a new session"
+            raise SessionNotActiveError(self._inactive_detail())
+
+    def _inactive_detail(self) -> Optional[str]:
+        """What Dataproc said about the inactive session, if the callback
+        knows (``is_active_callback.inactive_detail()``)."""
+        detail = getattr(self._is_active_callback, "inactive_detail", None)
+        if not callable(detail):
+            return None
+        try:
+            return detail()
+        except Exception:
+            logger.debug(
+                "Could not describe the inactive session", exc_info=True
             )
+            return None
 
     def _wrap_method(self, wrapped_method, method_name=None):
         if isinstance(method_name, bytes):

google/cloud/managed_spark_connect/exceptions.py

@@ -27,6 +27,28 @@ def _render_traceback_(self):
         return [self.message]
 
 
+class SessionNotActiveError(ManagedSparkConnectException, RuntimeError):
+    """The Managed Spark Session is gone: deleted, terminated or failed.
+
+    Raised by an RPC when Dataproc reports the session not found or in a
+    terminal state, and by ``getOrCreate()`` when a new session never became
+    usable. ``detail`` is what Dataproc said, e.g. ``"Managed Spark Session
+    sc-123 is FAILED: <reason>"``, when known.
+
+    A RuntimeError too, which is what this error used to be.
+    """
+
+    def __init__(self, detail=None):
+        self.detail = detail
+        message = "Session not active"
+        if detail:
+            message += f": {detail.rstrip('.')}"
+        super().__init__(f"{message}. Please create a new session")
+
+    def __reduce__(self):
+        return (type(self), (self.detail,))
+
+
 class PySparkNotInstalledError(ManagedSparkConnectException, ImportError):
     """Raised on import when pyspark is missing or incomplete.
 

google/cloud/managed_spark_connect/instrumentation/_grpc.py

@@ -70,6 +70,9 @@ def guarded() -> bool:
             _GUARD_NS.set(_GUARD_NS.get() + time.perf_counter_ns() - start)
 
     guarded._timed_guard = True  # type: ignore[attr-defined]
+    inactive_detail = getattr(callback, "inactive_detail", None)
+    if inactive_detail is not None:
+        guarded.inactive_detail = inactive_detail  # type: ignore[attr-defined]
     return guarded