diff --git a/README.md b/README.md index b5a7594..f16e677 100644 --- a/README.md +++ b/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. | diff --git a/docs/instrumentation.md b/docs/instrumentation.md index c4fd62e..71b6e0c 100644 --- a/docs/instrumentation.md +++ b/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.` (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: `), `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, diff --git a/google/cloud/managed_spark_connect/client/core.py b/google/cloud/managed_spark_connect/client/core.py index 88b138a..ebc8c70 100644 --- a/google/cloud/managed_spark_connect/client/core.py +++ b/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): diff --git a/google/cloud/managed_spark_connect/exceptions.py b/google/cloud/managed_spark_connect/exceptions.py index 16c5595..94b695e 100644 --- a/google/cloud/managed_spark_connect/exceptions.py +++ b/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: "``, 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. diff --git a/google/cloud/managed_spark_connect/instrumentation/_grpc.py b/google/cloud/managed_spark_connect/instrumentation/_grpc.py index f40a994..e8dd02f 100644 --- a/google/cloud/managed_spark_connect/instrumentation/_grpc.py +++ b/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 diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index 3b4406c..6eb3fe2 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -13,8 +13,9 @@ # limitations under the License. import atexit -import dataclasses from collections.abc import Mapping +import contextvars +import dataclasses import datetime import functools import html as html_lib @@ -70,7 +71,10 @@ install_release_coalescing, install_replay_on_position_lost, ) -from google.cloud.managed_spark_connect.exceptions import ManagedSparkConnectException +from google.cloud.managed_spark_connect.exceptions import ( + ManagedSparkConnectException, + SessionNotActiveError, +) from google.cloud.managed_spark_connect.instrumentation import ( NOOP_RECORDER, Instrumentation, @@ -82,6 +86,9 @@ install_client_hooks, uninstall_client_hooks, ) +from google.cloud.managed_spark_connect.instrumentation._core import ( + _CURRENT_SPAN, +) from google.cloud.managed_spark_connect.pypi_artifacts import PyPiArtifacts from google.cloud.dataproc_v1 import ( AuthenticationConfig, @@ -212,21 +219,28 @@ def _versions_text(spark_version: Optional[str]) -> str: return f"(Spark {spark_version})" if spark_version else "" -# Upper bound on how long the session messages wait for the server's Spark -# version (one AnalyzePlan RPC over the Spark Connect channel). On a newly -# created session this is the channel's first RPC, so it also warms up the -# proxy and channel (~4-5s when cold). +# Upper bound on how long the session-reuse messages wait for the server's +# Spark version (one AnalyzePlan RPC over the Spark Connect channel). _SPARK_VERSION_TIMEOUT_SECONDS = 10.0 +# How long getOrCreate() waits for a new session's first RPC (the Spark +# version) before asking Dataproc whether the session is still there. Long +# enough for a WebSocket handshake to time out (30s) and the liveness guard +# to ask Dataproc in turn. +_NEW_SESSION_CHECK_TIMEOUT_SECONDS = 60.0 -def _server_spark_version(session) -> Optional[str]: +def _server_spark_version( + session, timeout: Optional[float] = None +) -> Optional[str]: """The Spark Connect server's Spark version, or None. SparkSession.version on a Spark Connect session asks the server, so this costs one RPC; it runs on a daemon thread so a slow or failing server - delays the caller by at most _SPARK_VERSION_TIMEOUT_SECONDS. Never - raises. + delays the caller by at most ``timeout`` seconds (by default + _SPARK_VERSION_TIMEOUT_SECONDS). Never raises. """ + if timeout is None: + timeout = _SPARK_VERSION_TIMEOUT_SECONDS if session is None: return None result: Dict[str, Any] = {} @@ -237,12 +251,17 @@ def fetch(): except Exception: logger.debug("Could not get the Spark version", exc_info=True) + ctx = contextvars.copy_context() try: thread = threading.Thread( - target=fetch, name="spark-version", daemon=True + target=ctx.run, args=(fetch,), name="spark-version", daemon=True ) thread.start() - thread.join(_SPARK_VERSION_TIMEOUT_SECONDS) + thread.join(timeout) + if thread.is_alive(): + span = _CURRENT_SPAN.get() + if span is not None and span.name == "session.check_new": + span.set(timed_out=True) except Exception: return None version = result.get("version") @@ -258,6 +277,7 @@ def fetch(): SHOW_OPERATION_LINKS_ENV = "MANAGED_SPARK_CONNECT_SHOW_OPERATION_LINKS" PROGRESS_BAR_ENV = "MANAGED_SPARK_CONNECT_PROGRESS_BAR" PROGRESS_BAR_LEAVE_ENV = "MANAGED_SPARK_CONNECT_PROGRESS_BAR_LEAVE" +SESSION_START_RETRIES_ENV = "MANAGED_SPARK_CONNECT_SESSION_START_RETRIES" def _progress_bar_forced() -> bool: @@ -755,6 +775,50 @@ def __connect(self, session_response, session_name): return session + def __create_with_retries(self) -> "ManagedSparkSession": + """__create(), again with a new session ID when the new session + fails to start (SessionNotActiveError), up to + MANAGED_SPARK_CONNECT_SESSION_START_RETRIES times. + + Only for generated session IDs: a session the caller named + belongs to the caller, so its failure is theirs to see. + """ + retries = ( + 0 + if self._custom_session_id + else int( + environment.env_float( + SESSION_START_RETRIES_ENV, default=1, min_value=0 + ) + ) + ) + attempt = 1 + while True: + try: + return self.__create() + except SessionNotActiveError as exc: + retry = attempt <= retries + self._recorder.event( + "session.start_failed", + detail=exc.detail, + attempt=attempt, + retry=retry, + ) + if not retry: + raise + plain_message = ( + f"{exc.detail or 'Managed Spark Session failed to start'};" + " creating a new one" + ) + logger.warning( + "%s (retry %d of %d)", plain_message, attempt, retries + ) + self._output_element_or_message( + plain_message, + _status_line_html(html_lib.escape(plain_message)), + ) + attempt += 1 + def __create(self) -> "ManagedSparkSession": with self._lock: @@ -829,21 +893,6 @@ def create_session_pbar(): os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1" try: - if ( - os.getenv( - "MANAGED_SPARK_CONNECT_SESSION_TERMINATE_AT_EXIT", - "false", - ) - == "true" - ): - atexit.register( - lambda: terminate_s8s_session( - self._project_id, - self._region, - session_id, - self._client_options, - ) - ) with self._recorder.span("session.create_rpc"): operation = ( self.session_controller_client.create_session( @@ -911,19 +960,136 @@ def create_session_pbar(): ): create_session_pbar_thread.join() + try: + session_response = self._check_created_session( + session_response, session_config.name + ) + except SessionNotActiveError: + ManagedSparkSession._active_s8s_session_id = None + ManagedSparkSession._active_session_uses_custom_id = False + ManagedSparkSession._remove_stopped_session_from_file() + raise + except BaseException: + ManagedSparkSession._active_s8s_session_id = None + ManagedSparkSession._active_session_uses_custom_id = False + ManagedSparkSession._remove_stopped_session_from_file() + if not self._custom_session_id: + try: + terminate_s8s_session( + self._project_id, + self._region, + session_id, + self._client_options, + recorder=self._recorder, + ) + except Exception: + logger.debug( + "Failed to terminate session %s after" + " startup error", + session_id, + exc_info=True, + ) + raise + logger.debug( f"Managed Spark Session created: {session_id} in {int(time.time() - s8s_creation_start_time)} seconds" ) session = self._create_spark_connect_session_from_s8s( session_response, session_config.name ) - # After connecting, so the message can carry the server's - # Spark version (and warm the channel) in a single print. + try: + spark_version = self._check_new_session( + session, session_config.name + ) + except SessionNotActiveError: + session._abandon() + raise + except BaseException: + session._abandon(terminate=not self._custom_session_id) + raise + if ( + os.getenv( + "MANAGED_SPARK_CONNECT_SESSION_TERMINATE_AT_EXIT", + "false", + ) + == "true" + ): + atexit.register( + lambda: terminate_s8s_session( + self._project_id, + self._region, + session_id, + self._client_options, + ) + ) + # After the check, which fetched the server's Spark version + # (and warmed the channel) for the message. self._print_session_created_message( - self._session_versions(session) + _versions_text(spark_version) ) return session + def _check_new_session( + self, session: "ManagedSparkSession", session_name: str + ) -> Optional[str]: + """The new session's Spark version, once its first RPC got it. + + That RPC runs here, so a session that never became usable fails + getOrCreate() rather than the caller's first query. If it fails, + or takes more than _NEW_SESSION_CHECK_TIMEOUT_SECONDS, Dataproc + is asked about the session: one it reports gone raises + SessionNotActiveError saying why. Otherwise (including when + Dataproc cannot be asked) the session is returned as before, and + the caller's RPCs deal with whatever went wrong. + """ + with self._recorder.span("session.check_new") as span: + spark_version = _server_spark_version( + session, timeout=_NEW_SESSION_CHECK_TIMEOUT_SECONDS + ) + span.set(first_rpc_ok=spark_version is not None) + if spark_version is not None: + return spark_version + detail = _dead_session_detail( + session_name, + self._client_options, + recorder=self._recorder, + ) + span.set(dead=detail is not None) + if detail is None: + logger.debug( + "The first RPC to Managed Spark Session %s failed or" + " timed out, but Dataproc does not report it gone", + session_name, + ) + return None + raise SessionNotActiveError(detail) + + def _check_created_session( + self, session_response: Session, session_name: str + ) -> Session: + """The session CreateSession returned, once it can be connected to. + + As when reusing a session, that means ACTIVE with a Spark Connect + endpoint; otherwise this waits for one. A session that has already + failed, or is being terminated, raises SessionNotActiveError saying + what Dataproc reported, rather than leaving the first RPC to find + out. + """ + state = getattr(session_response, "state", None) + if state in _DEAD_STATES: + raise SessionNotActiveError( + _describe_inactive_session(session_name, session_response) + ) + if _has_spark_connect_endpoint(session_response): + return session_response + logger.debug( + "Managed Spark Session %s is %s without a Spark Connect" + " endpoint yet; waiting for one", + session_name, + getattr(state, "name", state), + ) + return self._wait_for_session_available(session_name) + def _wait_for_session_available( self, session_name: str, timeout: int = 1200 ) -> Session: @@ -946,21 +1112,22 @@ def _poll_for_session_endpoint( session = self.session_controller_client.get_session( name=session_name ) - if ( - "Spark Connect Server" in session.runtime_info.endpoints - and getattr(session, "state", None) - != Session.State.CREATING - ): - return session - if session.state in ( - Session.State.FAILED, - Session.State.TERMINATED, - ): - raise ManagedSparkConnectException( - f"Session {session_name} entered state {session.state.name} " - "before the Spark Connect endpoint became available." + # A dead session can still list its endpoint, so check + # the state first. + if session.state in _DEAD_STATES: + raise SessionNotActiveError( + f"{_describe_inactive_session(session_name, session)}" + " (before its Spark Connect endpoint became" + " available)" ) + if _has_spark_connect_endpoint(session): + return session time.sleep(5) + except NotFound as e: + raise SessionNotActiveError( + f"{_describe_inactive_session(session_name)}" + " (before its Spark Connect endpoint became available)" + ) from e except ManagedSparkConnectException: raise except Exception as e: @@ -1104,18 +1271,7 @@ def _get_exiting_active_session( if session_response is not None: self._last_session_response = session_response if session is None: - runtime_info = getattr( - session_response, "runtime_info", None - ) - endpoints = getattr(runtime_info, "endpoints", None) - has_endpoint = isinstance(endpoints, Mapping) and bool( - endpoints.get("Spark Connect Server") - ) - is_active_state = ( - getattr(session_response, "state", None) - == Session.State.ACTIVE - ) - if not (is_active_state and has_endpoint): + if not _has_spark_connect_endpoint(session_response): session_response = self._wait_for_session_available( session_name ) @@ -1216,7 +1372,7 @@ def getOrCreate(self) -> "ManagedSparkSession": mode = "reuse" if session is None: mode = "create" - session = self.__create() + session = self.__create_with_retries() span.set(mode=mode, **self._session_span_attrs()) finally: self._custom_session_response = None @@ -2212,6 +2368,22 @@ def stop(self, terminate: Optional[bool] = None) -> None: recorder.trace_id = "" self.instrumentation.close() + def _abandon(self, *, terminate: bool = False) -> None: + """Close a session that never became usable. + + By default (like stop(terminate=False)), the server session, which + Dataproc already reports gone (or is left to its TTL), is not + touched; pass terminate=True when abandoning a live generated-ID + session (e.g. on KeyboardInterrupt). Unlike stop(), the + instrumentation stays open, for getOrCreate() to go on recording + with it. + """ + try: + with NOOP_RECORDER.span("session.stop") as span: + self._stop(terminate, span) + finally: + get_recorder(self.instrumentation).trace_id = "" + def __del__(self) -> None: # pyspark's SparkSession.__del__ closes the client again. After # stop() the GC may run that on any thread, including a @@ -2363,7 +2535,11 @@ def terminate_s8s_session( terminate_session_request = TerminateSessionRequest() session_name = f"projects/{project_id}/locations/{region}/sessions/{active_s8s_session_id}" terminate_session_request.name = session_name - _record_liveness(_liveness_cache_key(session_name, client_options), False) + _record_liveness( + _liveness_cache_key(session_name, client_options), + False, + f"Managed Spark Session {active_s8s_session_id} was terminated", + ) state = None try: session_client = _create_session_controller_client( @@ -2416,10 +2592,12 @@ def terminate_s8s_session( class _LivenessEntry(NamedTuple): """The last definitive liveness verdict for a session and when it was - reached: ACTIVE, or dead (``NotFound`` or a terminal state).""" + reached: ACTIVE, or dead (``NotFound`` or a terminal state). ``detail`` + says why a dead session is dead, for its errors.""" active: bool at: float + detail: Optional[str] = None _LIVENESS_CACHE_LOCK = threading.Lock() @@ -2438,9 +2616,56 @@ def _liveness_cache_key( return (session_name, _get_client_option(client_options, "api_endpoint")) -def _record_liveness(key: tuple[str, Optional[str]], active: bool) -> None: +def _record_liveness( + key: tuple[str, Optional[str]], active: bool, detail: Optional[str] = None +) -> None: + with _LIVENESS_CACHE_LOCK: + _LIVENESS[key] = _LivenessEntry(active, _now(), detail) + + +_DEAD_STATES = ( + Session.State.FAILED, + Session.State.TERMINATING, + Session.State.TERMINATED, +) + + +def _has_spark_connect_endpoint(session: Optional[Session]) -> bool: + """Whether ``session`` is past CREATING and publishes a Spark Connect + endpoint.""" + runtime_info = getattr(session, "runtime_info", None) + endpoints = getattr(runtime_info, "endpoints", None) + return ( + getattr(session, "state", None) + in (Session.State.ACTIVE, Session.State.STATE_UNSPECIFIED) + and isinstance(endpoints, Mapping) + and bool(endpoints.get("Spark Connect Server")) + ) + + +def _describe_inactive_session( + session_name: str, session: Optional[Session] = None +) -> str: + """E.g. "Managed Spark Session sc-123 is FAILED: ", or "... was + deleted" without a ``session`` (``NotFound``).""" + subject = f"Managed Spark Session {session_name.rsplit('/', 1)[-1]}" + if session is None: + return f"{subject} was deleted" + state = getattr(session.state, "name", str(session.state)) + message = getattr(session, "state_message", "") + return f"{subject} is {state}" + (f": {message}" if message else "") + + +def _cached_inactive_detail( + session_name: str, client_options: Optional[ClientOptions] = None +) -> Optional[str]: + """Why the liveness cache holds ``session_name`` dead, if it does.""" + key = _liveness_cache_key(session_name, client_options) with _LIVENESS_CACHE_LOCK: - _LIVENESS[key] = _LivenessEntry(active, _now()) + entry = _LIVENESS.get(key) + if entry is None or entry.active: + return None + return entry.detail def _check_cached_liveness( @@ -2506,7 +2731,7 @@ def get_active_s8s_session_response( client_options=client_options ).get_session(get_session_request) except NotFound as e: - _record_liveness(key, False) + _record_liveness(key, False, _describe_inactive_session(session_name)) logger.debug("%s Managed Spark Session deleted: %s", session_name, e) return None state = get_session_response.state @@ -2515,10 +2740,56 @@ def get_active_s8s_session_response( return get_session_response if state == Session.State.CREATING: return get_session_response - _record_liveness(key, False) + _record_liveness( + key, + False, + _describe_inactive_session(session_name, get_session_response), + ) return None +def _dead_session_detail( + session_name: str, + client_options: Optional[ClientOptions] = None, + recorder: Any = NOOP_RECORDER, +) -> Optional[str]: + """What Dataproc says about a session it reports gone (``NotFound`` or a + terminal state), asking it now; None while the session is ACTIVE or + CREATING, or when Dataproc cannot be asked.""" + cached = _cached_inactive_detail(session_name, client_options) + if cached is not None: + return cached + key = _liveness_cache_key(session_name, client_options) + get_session_request = GetSessionRequest() + get_session_request.name = session_name + try: + with recorder.span( + "dataproc.get_session", reason="check_new", new_client=True + ): + response = _create_session_controller_client( + client_options=client_options + ).get_session(get_session_request) + except NotFound: + detail = _describe_inactive_session(session_name) + _record_liveness(key, False, detail) + return detail + except Exception: + logger.debug( + "Could not check Managed Spark Session %s", + session_name, + exc_info=True, + ) + return None + if response.state == Session.State.ACTIVE: + _record_liveness(key, True) + return None + if response.state == Session.State.CREATING: + return None + detail = _describe_inactive_session(session_name, response) + _record_liveness(key, False, detail) + return detail + + def is_s8s_session_active(session_name, client_options) -> bool: cached = _check_cached_liveness(session_name, client_options) if cached is not None: @@ -2528,6 +2799,28 @@ def is_s8s_session_active(session_name, client_options) -> bool: return True +class _IsActive: + """Callable session liveness guard with ``inactive_detail()`` for errors.""" + + __name__ = "is_active" + + def __init__( + self, + fn: Callable[[], bool], + session_name: str, + client_options: Any, + ): + self._fn = fn + self._session_name = session_name + self._client_options = client_options + + def __call__(self) -> bool: + return self._fn() + + def inactive_detail(self) -> Optional[str]: + return _cached_inactive_detail(self._session_name, self._client_options) + + def make_is_active_callback( session_name: str, client_options: Any, @@ -2577,4 +2870,4 @@ def is_active() -> bool: ) return active - return is_active + return _IsActive(is_active, session_name, client_options) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 1a675c0..c3adfb3 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -106,8 +106,10 @@ def dataproc_mocks(monkeypatch): return_value=True, ), mock.patch( + # Also a new session's first RPC: getting a version means the + # session is usable. "google.cloud.managed_spark_connect.session._server_spark_version", - return_value=None, + return_value="4.0.2", ), mock.patch( "google.cloud.managed_spark_connect.environment.get_client_environment_label", diff --git a/tests/unit/test_grpc_fork_noise.py b/tests/unit/test_grpc_fork_noise.py index 3bbcd13..d70ac51 100644 --- a/tests/unit/test_grpc_fork_noise.py +++ b/tests/unit/test_grpc_fork_noise.py @@ -318,7 +318,7 @@ def test_session_get_or_create_prewarms_and_stop_cleans_up_proxy(monkeypatch): ), mock.patch( "google.cloud.managed_spark_connect.session._server_spark_version", - return_value=None, + return_value="4.0.2", # The new session's first RPC worked. ), mock.patch("time.sleep"), ): diff --git a/tests/unit/test_guard.py b/tests/unit/test_guard.py index 80ff8b6..6e7d739 100644 --- a/tests/unit/test_guard.py +++ b/tests/unit/test_guard.py @@ -14,6 +14,7 @@ """Unit tests for the session liveness guard and TTL cache.""" import logging +import pickle from unittest import mock from google.api_core.client_options import ClientOptions @@ -34,10 +35,12 @@ from google.cloud.managed_spark_connect.client import core from google.cloud.managed_spark_connect.exceptions import ( ManagedSparkConnectException, + SessionNotActiveError, ) from google.cloud.managed_spark_connect.instrumentation import ( InMemorySink, Instrumentation, + get_recorder, ) SESSION_NAME = "projects/test-project/locations/us-central1/sessions/sess-1" @@ -103,6 +106,9 @@ def code(self): def trailing_metadata(self): return () + def result(self, timeout=None): + raise self + with ( mock.patch.object(core.proxy, "ManagedSparkSessionProxy") as proxy_cls, mock.patch.object(core, "DefaultChannelBuilder") as builder_cls, @@ -143,6 +149,91 @@ def trailing_metadata(self): assert failing_rpc.call_count == 1 +@pytest.mark.parametrize("instrumented", [False, True]) +@pytest.mark.parametrize( + "get_session, detail", + [ + ( + dict( + return_value=_make_session(Session.State.FAILED), + state_message="Insufficient capacity.", + ), + "Managed Spark Session sess-1 is FAILED: Insufficient capacity.", + ), + ( + dict(return_value=_make_session(Session.State.TERMINATED)), + "Managed Spark Session sess-1 is TERMINATED", + ), + ( + dict(side_effect=NotFound("gone")), + "Managed Spark Session sess-1 was deleted", + ), + ], +) +def test_not_active_error_says_what_dataproc_reported( + get_session, detail, instrumented +): + """The error names the session and its state (and why it failed), both + when the guard asks after a failed RPC and when a later RPC fast-fails + on the cached verdict.""" + + class FakeUnavailable(grpc.RpcError): + + def code(self): + return grpc.StatusCode.UNAVAILABLE + + def trailing_metadata(self): + return () + + def result(self, timeout=None): + raise self + + get_session = dict(get_session) + state_message = get_session.pop("state_message", None) + if state_message: + get_session["return_value"].state_message = state_message + inst = Instrumentation([InMemorySink()]) if instrumented else None + with ( + mock.patch.object(core.proxy, "ManagedSparkSessionProxy") as proxy_cls, + mock.patch.object(core, "DefaultChannelBuilder") as builder_cls, + mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client, + ): + proxy_cls.return_value.port = 12345 + fake_inner = mock.MagicMock() + failing = mock.MagicMock( + side_effect=FakeUnavailable("handshake timed out") + ) + failing.with_call.side_effect = FakeUnavailable("handshake timed out") + fake_inner.unary_unary.return_value = failing + builder_cls.return_value.toChannel.return_value = fake_inner + create_client.return_value.get_session.configure_mock(**get_session) + + cb = session_mod.make_is_active_callback( + SESSION_NAME, + CLIENT_OPTIONS, + recorder=get_recorder(inst), + ) + method = core.ProxiedChannel( + "example.com", cb, instrumentation=inst + ).unary_unary("/spark.connect.SparkConnectService/Config") + for request in ("after-error", "cached"): + with pytest.raises(SessionNotActiveError) as excinfo: + method(request) + assert excinfo.value.detail == detail + assert str(excinfo.value) == ( + f"Session not active: {detail.rstrip('.')}. Please create a" + " new session" + ) + roundtripped = pickle.loads(pickle.dumps(excinfo.value)) + assert roundtripped.detail == detail + assert str(roundtripped) == str(excinfo.value) + # Still what callers caught before. + assert isinstance(excinfo.value, RuntimeError) + assert isinstance(excinfo.value, ManagedSparkConnectException) + + def test_lazy_mode_reraises_rpc_error_when_session_is_still_active(): """On RPC error, if GetSession reports session ACTIVE, the original RpcError is re-raised.""" diff --git a/tests/unit/test_init.py b/tests/unit/test_init.py index 7118158..12a1dad 100644 --- a/tests/unit/test_init.py +++ b/tests/unit/test_init.py @@ -1142,6 +1142,8 @@ def test_all_import_orders_load_all_extras(self): mock.patch("google.cloud.dataproc_v1.SessionControllerClient") as mock_client_cls, mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") as mock_config, mock.patch("google.cloud.managed_spark_connect.session.is_s8s_session_active", return_value=True), + # The new session's first RPC, which would go to the fake server. + mock.patch("google.cloud.managed_spark_connect.session._server_spark_version", return_value="4.0.2"), ): cred = mock.MagicMock() cred.token = "token" diff --git a/tests/unit/test_instrumentation.py b/tests/unit/test_instrumentation.py index c0a445a..4c87834 100644 --- a/tests/unit/test_instrumentation.py +++ b/tests/unit/test_instrumentation.py @@ -731,7 +731,7 @@ def test_guard_failure_is_an_error_on_the_outer_span(proxied): rec = one(inst, "grpc.Config") assert rec["attrs"]["layer"] == "outer" assert rec["status"] == "error" - assert rec["error"]["type"] == "RuntimeError" + assert rec["error"]["type"] == "SessionNotActiveError" assert rec["attrs"]["guard_ns"] >= 0 assert spans(inst, "grpc.Config") == [rec] # inner never ran (emitted,) = [r for r in snapshots if r.get("name") == "grpc.Config"] diff --git a/tests/unit/test_logging_hygiene.py b/tests/unit/test_logging_hygiene.py index 7499fce..1cadb5c 100644 --- a/tests/unit/test_logging_hygiene.py +++ b/tests/unit/test_logging_hygiene.py @@ -296,6 +296,11 @@ def test_non_interactive_lifecycle_zero_stdout_stderr_and_warnings( mock.patch.object( ManagedSparkSession.Builder, "_check_pyspark_version_compatibility" ), + mock.patch( + # The new session's first RPC, which would go to the fake server. + "google.cloud.managed_spark_connect.session._server_spark_version", + return_value="4.0.2", + ), mock.patch("tqdm.trange") as mock_trange, warnings.catch_warnings(record=True) as caught_warnings, ): diff --git a/tests/unit/test_polling.py b/tests/unit/test_polling.py index 5e640b4..c5585cd 100644 --- a/tests/unit/test_polling.py +++ b/tests/unit/test_polling.py @@ -95,7 +95,7 @@ def test_terminate_s8s_session_sleeps_only_between_active_polls( @mock.patch( "google.cloud.managed_spark_connect.session._server_spark_version", - return_value=None, + return_value="4.0.2", # The new session's first RPC worked. ) @mock.patch("google.auth.default") @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") diff --git a/tests/unit/test_reattach_patch.py b/tests/unit/test_reattach_patch.py index 5899f99..d808ef4 100644 --- a/tests/unit/test_reattach_patch.py +++ b/tests/unit/test_reattach_patch.py @@ -461,6 +461,11 @@ def fake_fetch(self, req, observations, self_destruct=False): "google.cloud.managed_spark_connect.session.is_s8s_session_active", return_value=True, ), + mock.patch( + # The new session's first RPC worked. + "google.cloud.managed_spark_connect.session._server_spark_version", + return_value="4.0.2", + ), mock.patch("time.sleep"), ): op = mock.Mock() diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index d6ea70c..a379e11 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -73,10 +73,11 @@ def setUp(self): ) patcher.start() self.addCleanup(patcher.stop) - # Never ask a (fake) Spark Connect server for its version. + # Never ask a (fake) Spark Connect server for its version. Getting + # one is also what tells getOrCreate() a new session is usable. patcher = mock.patch( "google.cloud.managed_spark_connect.session._server_spark_version", - return_value=None, + return_value="4.0.2", ) self.mock_server_spark_version = patcher.start() self.addCleanup(patcher.stop) @@ -1585,6 +1586,7 @@ def _get_existing_session_output(self): def test_get_exiting_active_session_prints_only_text_links(self): self._patch_ide(self._AGY_IDE) + self.mock_server_spark_version.return_value = None printed = self._get_existing_session_output() @@ -1798,6 +1800,7 @@ def test_session_versions(self): "google.cloud.managed_spark_connect.environment.is_interactive", return_value=True, ): + self.mock_server_spark_version.return_value = None self.assertEqual(builder._session_versions(None), "") self.mock_server_spark_version.return_value = "4.0.2" self.assertEqual( @@ -2071,10 +2074,11 @@ def setUp(self): ) patcher.start() self.addCleanup(patcher.stop) - # Never ask a (fake) Spark Connect server for its version. + # Never ask a (fake) Spark Connect server for its version. Getting + # one is also what tells getOrCreate() a new session is usable. patcher = mock.patch( "google.cloud.managed_spark_connect.session._server_spark_version", - return_value=None, + return_value="4.0.2", ) self.mock_server_spark_version = patcher.start() self.addCleanup(patcher.stop) @@ -3585,6 +3589,11 @@ def test_session_terminating_waits_then_deletes_and_warns( self.assertIn("TERMINATING", warning_msg) @pytest.mark.slow + @mock.patch( + # The recreated session's first RPC worked. + "google.cloud.managed_spark_connect.session._server_spark_version", + return_value="4.0.2", + ) @mock.patch("google.auth.default") @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @@ -3599,6 +3608,7 @@ def test_get_or_create_with_terminated_custom_id_deletes_and_recreates( mock_client_config, mock_dataproc_session_controller_client, mock_credentials, + _mock_server_spark_version, ): """Test that getOrCreate deletes a terminated custom-ID session and creates a new one.""" mock_is_s8s_session_active.return_value = True diff --git a/tests/unit/test_session_start.py b/tests/unit/test_session_start.py new file mode 100644 index 0000000..9812c60 --- /dev/null +++ b/tests/unit/test_session_start.py @@ -0,0 +1,549 @@ +# Copyright 2026 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. +"""getOrCreate() on a new session that does not come up.""" + +from pathlib import Path +import threading +from unittest import mock + +from google.api_core.exceptions import NotFound, ServiceUnavailable +from google.cloud.dataproc_v1 import Session +import pytest + +from google.cloud.managed_spark_connect import ManagedSparkSession +from google.cloud.managed_spark_connect import session as session_mod +from google.cloud.managed_spark_connect.exceptions import SessionNotActiveError +from google.cloud.managed_spark_connect.instrumentation import ( + InMemorySink, + Instrumentation, + get_recorder, +) + +from .conftest import session_response + +_REAL_SERVER_SPARK_VERSION = session_mod._server_spark_version + + +def _builder(): + return ManagedSparkSession.builder.projectId("test-project").location( + "test-region" + ) + + +def _failed(session_id, uuid, state=Session.State.FAILED, message=""): + response = session_response(session_id, uuid) + response.state = state + response.state_message = message + return response + + +@pytest.fixture +def no_retries(dataproc_mocks, monkeypatch): + """One attempt only, for the tests of how a single new session fails.""" + monkeypatch.setenv(session_mod.SESSION_START_RETRIES_ENV, "0") + return dataproc_mocks + + +@pytest.mark.parametrize( + "state, message, detail", + [ + ( + Session.State.FAILED, + "Insufficient capacity", + "is FAILED: Insufficient capacity", + ), + (Session.State.TERMINATING, "", "is TERMINATING"), + ( + Session.State.TERMINATED, + "Idle TTL expired.", + "is TERMINATED: Idle TTL expired.", + ), + ], +) +def test_create_result_already_dead_raises_with_dataproc_state( + no_retries, state, message, detail +): + """A session CreateSession returns failed (or terminating) is not + connected to: getOrCreate() says what Dataproc reported and leaves no + session behind.""" + mocks = no_retries + mocks["client"].create_session.return_value.result.return_value = _failed( + mocks["session_id"], mocks["uuid"], state, message + ) + + with pytest.raises(SessionNotActiveError) as excinfo: + _builder().getOrCreate() + + assert excinfo.value.detail == ( + f"Managed Spark Session {mocks['session_id']} {detail}" + ) + assert ".." not in str(excinfo.value) + assert ManagedSparkSession._active_s8s_session_id is None + assert ManagedSparkSession._default_session is None + + +@pytest.mark.parametrize( + "initial_state, endpoints", + [ + (Session.State.CREATING, {}), + (Session.State.ACTIVE, {}), + (Session.State.ACTIVE, {"Spark Connect Server": ""}), + ], +) +def test_create_result_without_endpoint_waits_for_one( + dataproc_mocks, initial_state, endpoints +): + """A session that CreateSession returns before its Spark Connect + endpoint is published is polled until it has one, as when reusing a + session.""" + mocks = dataproc_mocks + creating = session_response(mocks["session_id"], mocks["uuid"]) + creating.state = initial_state + creating.runtime_info.endpoints = endpoints + mocks["client"].create_session.return_value.result.return_value = creating + ready = session_response(mocks["session_id"], mocks["uuid"]) + mocks["client"].get_session.side_effect = [creating, ready] + + session = _builder().getOrCreate() + try: + assert mocks["client"].get_session.call_count == 2 + assert ManagedSparkSession._active_s8s_session_uuid == mocks["uuid"] + finally: + mocks["client"].get_session.side_effect = None + session.stop() + + +def test_waiting_for_endpoint_gives_up_when_the_session_fails(no_retries): + mocks = no_retries + creating = session_response(mocks["session_id"], mocks["uuid"]) + creating.state = Session.State.CREATING + creating.runtime_info.endpoints = {} + mocks["client"].create_session.return_value.result.return_value = creating + mocks["client"].get_session.side_effect = [ + creating, + _failed(mocks["session_id"], mocks["uuid"], message="Out of quota"), + ] + + with pytest.raises(SessionNotActiveError, match="is FAILED: Out of quota"): + _builder().getOrCreate() + assert ManagedSparkSession._active_s8s_session_id is None + + +def test_waiting_for_endpoint_gives_up_when_the_session_is_deleted(no_retries): + mocks = no_retries + creating = session_response(mocks["session_id"], mocks["uuid"]) + creating.state = Session.State.CREATING + creating.runtime_info.endpoints = {} + mocks["client"].create_session.return_value.result.return_value = creating + mocks["client"].get_session.side_effect = NotFound("gone") + + with pytest.raises(SessionNotActiveError) as excinfo: + _builder().getOrCreate() + + assert excinfo.value.detail == ( + f"Managed Spark Session {mocks['session_id']} was deleted" + " (before its Spark Connect endpoint became available)" + ) + assert mocks["client"].get_session.call_count == 1 + assert ManagedSparkSession._active_s8s_session_id is None + + +def test_waiting_for_endpoint_timeout_terminates_generated_session_without_retry( + dataproc_mocks, two_session_ids +): + mocks = dataproc_mocks + creating = session_response("sc-first", "uuid-1") + creating.state = Session.State.CREATING + creating.runtime_info.endpoints = {} + mocks["client"].create_session.return_value.result.return_value = creating + + with mock.patch.object( + ManagedSparkSession.Builder, + "_poll_for_session_endpoint", + side_effect=RuntimeError( + "Timed out waiting for Spark Connect endpoint" + ), + ): + with pytest.raises( + RuntimeError, match="Timed out waiting for Spark Connect endpoint" + ): + _builder().getOrCreate() + + assert _created_ids(mocks) == ["sc-first"] + (terminated,) = mocks["client"].terminate_session.call_args_list + assert terminated.args[0].name.endswith("/sc-first") + assert ManagedSparkSession._active_s8s_session_id is None + + +@pytest.fixture +def first_rpc_fails(): + """The new session's first RPC (the Spark version) fails, as when its + WebSocket handshake times out.""" + with mock.patch( + "google.cloud.managed_spark_connect.session._server_spark_version", + return_value=None, + ) as spark_version: + yield spark_version + + +def test_unusable_new_session_fails_get_or_create(no_retries, first_rpc_fails): + """The CI failure mode: CreateSession succeeds, the first RPC fails and + Dataproc then reports the session FAILED. getOrCreate() raises saying + so, instead of returning a session whose first query fails with a bare + "Session not active", and leaves nothing behind.""" + mocks = no_retries + mocks["client"].get_session.return_value = _failed( + mocks["session_id"], mocks["uuid"], message="Agent failed to start" + ) + + with pytest.raises(SessionNotActiveError) as excinfo: + _builder().getOrCreate() + + assert excinfo.value.detail == ( + f"Managed Spark Session {mocks['session_id']} is FAILED: Agent failed" + " to start" + ) + (call,) = first_rpc_fails.call_args_list + assert ( + call.kwargs["timeout"] == session_mod._NEW_SESSION_CHECK_TIMEOUT_SECONDS + ) + abandoned = call.args[0] + assert abandoned.client.is_closed + assert ManagedSparkSession._active_s8s_session_id is None + assert ManagedSparkSession._default_session is None + # The session is dead already; nothing tries to terminate it. + mocks["client"].terminate_session.assert_not_called() + + +@pytest.mark.parametrize( + "get_session", + [ + pytest.param("active", id="dataproc-reports-active"), + pytest.param(ServiceUnavailable("down"), id="dataproc-unreachable"), + ], +) +def test_new_session_is_returned_unless_dataproc_reports_it_gone( + dataproc_mocks, first_rpc_fails, get_session +): + """A failed or slow first RPC alone does not fail getOrCreate(): only + Dataproc reporting the session gone does.""" + mocks = dataproc_mocks + stop_response = mocks["client"].get_session.return_value + if get_session == "active": + mocks["client"].get_session.return_value = mocks["response"] + else: + mocks["client"].get_session.side_effect = get_session + + session = _builder().getOrCreate() + try: + assert ManagedSparkSession._default_session is session + assert not session.client.is_closed + finally: + mocks["client"].get_session.side_effect = None + mocks["client"].get_session.return_value = stop_response + session.stop() + + +def test_new_session_check_is_recorded(no_retries, first_rpc_fails): + mocks = no_retries + mocks["client"].get_session.return_value = _failed( + mocks["session_id"], mocks["uuid"] + ) + inst = Instrumentation([InMemorySink()]) + + with pytest.raises(SessionNotActiveError): + _builder().instrumentation(inst).getOrCreate() + + (span,) = [r for r in inst.records if r.get("name") == "session.check_new"] + assert span["attrs"]["first_rpc_ok"] is False + assert span["attrs"]["dead"] is True + assert span["status"] == "error" + (get_span,) = [ + r for r in inst.records if r.get("name") == "dataproc.get_session" + ] + assert get_span["attrs"]["reason"] == "check_new" + assert get_span["parent_id"] == span["span_id"] + + +def test_server_spark_version_propagates_span_context_and_reuses_guard_liveness( + no_retries, +): + """When _server_spark_version runs on its worker thread, contextvars are + copied so spans inside session.version parent under session.check_new, and + if the channel guard already cached the session's inactive detail on the + failed RPC, _dead_session_detail reuses it without a second GetSession.""" + mocks = no_retries + mocks["client"].get_session.return_value = _failed( + mocks["session_id"], mocks["uuid"], message="Agent crashed" + ) + inst = Instrumentation([InMemorySink()]) + session_name = f"projects/test-project/locations/test-region/sessions/{mocks['session_id']}" + + def fake_version(self): + recorder = get_recorder(self.instrumentation) + with recorder.span("grpc.AnalyzePlan"): + # Simulate the channel guard's on_error running make_is_active_callback + # when the AnalyzePlan RPC for session.version fails. + cb = session_mod.make_is_active_callback(session_name, None) + assert cb() is False + raise SessionNotActiveError(cb.inactive_detail()) + + with ( + mock.patch( + "google.cloud.managed_spark_connect.session._server_spark_version", + side_effect=_REAL_SERVER_SPARK_VERSION, + ), + mock.patch.object( + ManagedSparkSession, + "version", + property(fake_version), + ), + ): + with pytest.raises(SessionNotActiveError) as excinfo: + _builder().instrumentation(inst).getOrCreate() + + assert excinfo.value.detail == ( + f"Managed Spark Session {mocks['session_id']} is FAILED: Agent crashed" + ) + # Only the guard's GetSession ran; _dead_session_detail reused _LIVENESS. + assert mocks["client"].get_session.call_count == 1 + (check_span,) = [ + r for r in inst.records if r.get("name") == "session.check_new" + ] + (analyze_span,) = [ + r for r in inst.records if r.get("name") == "grpc.AnalyzePlan" + ] + assert analyze_span["parent_id"] == check_span["span_id"] + assert "timed_out" not in check_span["attrs"] + + +def test_server_spark_version_timeout_sets_timed_out_attr(): + inst = Instrumentation([InMemorySink()]) + recorder = get_recorder(inst) + unblock = threading.Event() + + class SlowSession: + + @property + def version(self): + unblock.wait(timeout=5.0) + return "4.0.2" + + try: + with recorder.span("session.check_new"): + result = session_mod._server_spark_version( + SlowSession(), timeout=0.01 + ) + finally: + unblock.set() + + assert result is None + (check_span,) = [ + r for r in inst.records if r.get("name") == "session.check_new" + ] + assert check_span["attrs"]["timed_out"] is True + + +# --------------------------------------------------------------- retrying + + +@pytest.fixture +def two_session_ids(dataproc_mocks): + """Each attempt creates a session under a new ID, as generated IDs do.""" + with mock.patch.object( + ManagedSparkSession.Builder, + "generate_session_id", + side_effect=["sc-first", "sc-second"], + ): + yield "sc-first", "sc-second" + + +def _get_session_failed_then_terminating(mocks, message): + """GetSession: FAILED for the check of the first session; TERMINATING + from then on, which is what stop() waits for.""" + calls = [] + + def get_session(request, *args, **kwargs): + calls.append(request) + if len(calls) == 1: + return _failed("sc-first", "uuid-1", message=message) + return mocks["client"].get_session.return_value + + mocks["client"].get_session.side_effect = get_session + + +def _created_ids(mocks): + return [ + call.args[0].session_id + for call in mocks["client"].create_session.call_args_list + ] + + +def test_session_that_fails_to_start_is_replaced( + dataproc_mocks, two_session_ids, caplog +): + """The CI failure mode, recovered: the first new session's first RPC + fails and Dataproc reports it FAILED, so getOrCreate() creates another + one (a new ID) and returns that.""" + mocks = dataproc_mocks + mocks["client"].create_session.return_value.result.side_effect = [ + session_response("sc-first", "11111111-1111-1111-1111-111111111111"), + session_response("sc-second", "22222222-2222-2222-2222-222222222222"), + ] + _get_session_failed_then_terminating(mocks, "Agent failed to start") + inst = Instrumentation([InMemorySink()]) + + with ( + caplog.at_level("WARNING", logger="google.cloud.managed_spark_connect"), + mock.patch( + "google.cloud.managed_spark_connect.session._server_spark_version", + side_effect=[None, "4.0.2"], + ) as spark_version, + ): + session = _builder().instrumentation(inst).getOrCreate() + try: + assert _created_ids(mocks) == ["sc-first", "sc-second"] + first, second = [call.args[0] for call in spark_version.call_args_list] + assert first.client.is_closed + assert second is session + assert ManagedSparkSession._default_session is session + assert ManagedSparkSession._active_s8s_session_id == "sc-second" + assert ( + "Managed Spark Session sc-first is FAILED: Agent failed to start;" + " creating a new one (retry 1 of 1)" in caplog.text + ) + (event,) = [ + r for r in inst.records if r.get("name") == "session.start_failed" + ] + assert event["attrs"] == { + "detail": "Managed Spark Session sc-first is FAILED: Agent failed" + " to start", + "attempt": 1, + "retry": True, + } + # Abandoning the failed first attempt does not emit a session.stop + # span inside session.get_or_create, and clears Attempt 1's trace_id + # before Attempt 2 starts. + assert [ + r for r in inst.records if r.get("name") == "session.stop" + ] == [] + create_spans = [ + r for r in inst.records if r.get("name") == "session.create_rpc" + ] + assert [s["trace_id"] for s in create_spans] == ["", ""] + (get_or_create_span,) = [ + r for r in inst.records if r.get("name") == "session.get_or_create" + ] + assert ( + get_or_create_span["trace_id"] + == "22222222-2222-2222-2222-222222222222" + ) + # The first session is already dead: only the second is terminated. + finally: + mocks["client"].get_session.side_effect = None + session.stop() + (terminated,) = mocks["client"].terminate_session.call_args_list + assert terminated.args[0].name.endswith("/sc-second") + + +def test_active_session_file_and_atexit_cleaned_up_on_failed_start( + dataproc_mocks, two_session_ids, monkeypatch, tmp_path: Path +): + mocks = dataproc_mocks + active_file = tmp_path / "active_session.json" + monkeypatch.setenv( + "MANAGED_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH", str(active_file) + ) + monkeypatch.setenv( + "MANAGED_SPARK_CONNECT_SESSION_TERMINATE_AT_EXIT", "true" + ) + monkeypatch.setenv(session_mod.SESSION_START_RETRIES_ENV, "0") + mocks["client"].create_session.return_value.result.return_value = _failed( + "sc-first", "uuid-1", message="Out of quota" + ) + + with mock.patch("atexit.register") as mock_atexit: + with pytest.raises(SessionNotActiveError): + _builder().getOrCreate() + + assert active_file.read_text() == "" + mock_atexit.assert_not_called() + + +def test_session_already_dead_at_create_is_replaced( + dataproc_mocks, two_session_ids +): + mocks = dataproc_mocks + mocks["client"].create_session.return_value.result.side_effect = [ + _failed("sc-first", "uuid-1", message="Out of quota"), + session_response("sc-second", "22222222-2222-2222-2222-222222222222"), + ] + + session = _builder().getOrCreate() + try: + assert _created_ids(mocks) == ["sc-first", "sc-second"] + assert ManagedSparkSession._active_s8s_session_uuid == ( + "22222222-2222-2222-2222-222222222222" + ) + finally: + session.stop() + + +def test_gives_up_after_the_configured_retries( + dataproc_mocks, two_session_ids, monkeypatch +): + mocks = dataproc_mocks + mocks["client"].create_session.return_value.result.side_effect = [ + _failed("sc-first", "uuid-1", message="Out of quota"), + _failed("sc-second", "uuid-2", message="Still out of quota"), + ] + inst = Instrumentation([InMemorySink()]) + + with pytest.raises(SessionNotActiveError, match="Still out of quota"): + _builder().instrumentation(inst).getOrCreate() + + assert _created_ids(mocks) == ["sc-first", "sc-second"] + events = [ + r for r in inst.records if r.get("name") == "session.start_failed" + ] + assert [(e["attrs"]["attempt"], e["attrs"]["retry"]) for e in events] == [ + (1, True), + (2, False), + ] + assert ManagedSparkSession._active_s8s_session_id is None + + +def test_retries_can_be_turned_off(no_retries, two_session_ids): + mocks = no_retries + mocks["client"].create_session.return_value.result.return_value = _failed( + "sc-first", "uuid-1" + ) + + with pytest.raises(SessionNotActiveError): + _builder().getOrCreate() + assert _created_ids(mocks) == ["sc-first"] + + +def test_named_session_that_fails_to_start_is_not_replaced(dataproc_mocks): + """A session the caller named is theirs: its failure is reported, not + papered over with another session.""" + mocks = dataproc_mocks + mocks["client"].get_session.side_effect = NotFound("no such session") + mocks["client"].create_session.return_value.result.return_value = _failed( + "my-session", "uuid-1", message="Out of quota" + ) + + with pytest.raises(SessionNotActiveError, match="my-session is FAILED"): + _builder().sessionId("my-session").getOrCreate() + assert _created_ids(mocks) == ["my-session"]