From ba7c9166dca223308094dbee0d5852c6e0dd7a1b Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Fri, 30 Oct 2026 15:11:04 -0700 Subject: [PATCH 1/4] feat(session): say why a session is not active An RPC on a session that Dataproc reports gone used to fail with a bare RuntimeError("Session not active. Please create a new session"), dropping the state the liveness check had just fetched. It now raises SessionNotActiveError, which names the session and what Dataproc said: "Managed Spark Session sc-123 is FAILED: ", "... is TERMINATED", "... was deleted". The liveness cache keeps that detail with a dead verdict, so RPCs that fast-fail on the cached verdict say the same. SessionNotActiveError subclasses RuntimeError, as before, and ManagedSparkConnectException, and its message keeps the old prefix and suffix. --- README.md | 2 +- .../managed_spark_connect/client/core.py | 17 +++- .../cloud/managed_spark_connect/exceptions.py | 22 +++++ .../instrumentation/_grpc.py | 3 + google/cloud/managed_spark_connect/session.py | 73 +++++++++++++-- tests/unit/test_guard.py | 91 +++++++++++++++++++ tests/unit/test_instrumentation.py | 2 +- 7 files changed, 199 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index b5a75940..04b769e2 100644 --- a/README.md +++ b/README.md @@ -138,7 +138,7 @@ 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. | | `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/google/cloud/managed_spark_connect/client/core.py b/google/cloud/managed_spark_connect/client/core.py index 88b138ad..ebc8c705 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 16c55952..94b695ed 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 f40a9945..e8dd02f6 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 3b4406cb..fc739964 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -2363,7 +2363,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 +2420,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 +2444,36 @@ 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()) + _LIVENESS[key] = _LivenessEntry(active, _now(), detail) + + +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: + entry = _LIVENESS.get(key) + if entry is None or entry.active: + return None + return entry.detail def _check_cached_liveness( @@ -2506,7 +2539,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,7 +2548,11 @@ 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 @@ -2528,6 +2565,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 +2636,4 @@ def is_active() -> bool: ) return active - return is_active + return _IsActive(is_active, session_name, client_options) diff --git a/tests/unit/test_guard.py b/tests/unit/test_guard.py index 80ff8b62..6e7d7396 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_instrumentation.py b/tests/unit/test_instrumentation.py index c0a445ad..4c878344 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"] From bfe3ee059afd3178447b07d158cc019c7db4148f Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Fri, 30 Oct 2026 15:12:12 -0700 Subject: [PATCH 2/4] fix(session): check the session CreateSession returns before connecting getOrCreate() connected to whatever session the CreateSession operation returned, without looking at its state or its Spark Connect endpoint. Reusing a session already checks both; a new one now does too: a session that is already FAILED, TERMINATING or TERMINATED raises SessionNotActiveError with what Dataproc reported, and one that is still CREATING or has no endpoint yet is polled until it has one. Polling for the endpoint also checked the endpoint before the state, so it returned a session that had failed after publishing its endpoint. It now checks for FAILED, TERMINATING (new) or TERMINATED first, and raises SessionNotActiveError with the state message instead of a bare ManagedSparkConnectException (which SessionNotActiveError subclasses). --- google/cloud/managed_spark_connect/session.py | 122 ++++++++++++++---- tests/unit/test_session_start.py | 104 +++++++++++++++ 2 files changed, 200 insertions(+), 26 deletions(-) create mode 100644 tests/unit/test_session_start.py diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index fc739964..5a221364 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -70,7 +70,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, @@ -911,6 +914,37 @@ 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" ) @@ -924,6 +958,32 @@ def create_session_pbar(): ) return session + 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 +1006,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 +1165,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 ) @@ -2451,6 +2501,26 @@ def _record_liveness( _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: diff --git a/tests/unit/test_session_start.py b/tests/unit/test_session_start.py new file mode 100644 index 00000000..4a1715f3 --- /dev/null +++ b/tests/unit/test_session_start.py @@ -0,0 +1,104 @@ +# 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 google.cloud.dataproc_v1 import Session +import pytest + +from google.cloud.managed_spark_connect import ManagedSparkSession +from google.cloud.managed_spark_connect.exceptions import SessionNotActiveError + +from .conftest import session_response + + +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.mark.parametrize( + "state, message, detail", + [ + ( + Session.State.FAILED, + "Insufficient capacity", + "is FAILED: Insufficient capacity", + ), + (Session.State.TERMINATING, "", "is TERMINATING"), + ], +) +def test_create_result_already_dead_raises_with_dataproc_state( + dataproc_mocks, 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 = dataproc_mocks + 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 ManagedSparkSession._active_s8s_session_id is None + assert ManagedSparkSession._default_session is None + + +def test_create_result_without_endpoint_waits_for_one(dataproc_mocks): + """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 = Session.State.CREATING + creating.runtime_info.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(dataproc_mocks): + mocks = dataproc_mocks + 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 From c12084629e9b45b108f447e06a7abe704a8347ef Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Fri, 30 Oct 2026 15:33:50 -0700 Subject: [PATCH 3/4] feat(session): check that a new session works before getOrCreate() returns getOrCreate() returned a new session without any RPC to it outside notebooks, so a session that came up broken was only found out by the caller's first query: in CI, an 11-minute CreateSession, then a WebSocket handshake timeout and "Session not active" on the first createDataFrame(). getOrCreate() now makes the new session's first RPC itself: the Spark version, which the session-created message already fetched in notebooks and now gets from here. If it fails, or takes more than 60s (enough for a 30s WebSocket handshake timeout plus the liveness guard), Dataproc is asked about the session; one it reports gone is closed client-side and getOrCreate() raises SessionNotActiveError saying why. Anything else, including Dataproc being unreachable, returns the session as before. The check is recorded as a session.check_new span (first_rpc_ok, dead, timed_out). Tests that mock the Spark version now return one, since a version is what tells getOrCreate() that the new session is usable. --- docs/instrumentation.md | 4 +- google/cloud/managed_spark_connect/session.py | 145 ++++++++++++++++-- tests/unit/conftest.py | 4 +- tests/unit/test_grpc_fork_noise.py | 2 +- tests/unit/test_init.py | 2 + tests/unit/test_logging_hygiene.py | 5 + tests/unit/test_polling.py | 2 +- tests/unit/test_reattach_patch.py | 5 + tests/unit/test_session.py | 18 ++- tests/unit/test_session_start.py | 95 ++++++++++++ 10 files changed, 260 insertions(+), 22 deletions(-) diff --git a/docs/instrumentation.md b/docs/instrumentation.md index c4fd62e7..ef5ec5f1 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` | diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index 5a221364..b45a3be3 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 @@ -85,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, @@ -215,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] = {} @@ -240,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") @@ -951,13 +967,58 @@ def create_session_pbar(): 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 + # 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: @@ -2262,6 +2323,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 @@ -2626,6 +2703,48 @@ def get_active_s8s_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: diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 1a675c07..c3adfb38 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 3bbcd133..d70ac518 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_init.py b/tests/unit/test_init.py index 71181589..12a1dad9 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_logging_hygiene.py b/tests/unit/test_logging_hygiene.py index 7499fceb..1cadb5c3 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 5e640b4b..c5585cdc 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 5899f99e..d808ef45 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 d6ea70cf..a379e115 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 index 4a1715f3..613c824b 100644 --- a/tests/unit/test_session_start.py +++ b/tests/unit/test_session_start.py @@ -13,11 +13,19 @@ # limitations under the License. """getOrCreate() on a new session that does not come up.""" +from unittest import mock + +from google.api_core.exceptions import 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, +) from .conftest import session_response @@ -102,3 +110,90 @@ def test_waiting_for_endpoint_gives_up_when_the_session_fails(dataproc_mocks): with pytest.raises(SessionNotActiveError, match="is FAILED: Out of quota"): _builder().getOrCreate() 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( + dataproc_mocks, 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 = dataproc_mocks + 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(dataproc_mocks, first_rpc_fails): + mocks = dataproc_mocks + 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" From 958c02b79305f22b387599cd12e9e311d2a80764 Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Fri, 30 Oct 2026 15:46:59 -0700 Subject: [PATCH 4/4] feat(session): replace a new session that fails to start When a session getOrCreate() just created turns out dead before it could be used (Dataproc reports it failed or terminated in the CreateSession result, while waiting for its endpoint, or after its first RPC failed), getOrCreate() now creates another session in its place and returns that, up to MANAGED_SPARK_CONNECT_SESSION_START_RETRIES times (default 1, 0 to disable). Nothing has run on the dead session, so nothing is lost but time; the dead session itself is left alone, since Dataproc already has it failed or terminated. Only sessions with a generated ID are replaced: a session created with sessionId() belongs to the caller, so its failure still raises SessionNotActiveError. A replacement logs a WARNING, shows a status line in notebooks, and is recorded as a session.start_failed event (detail, attempt, retry). --- README.md | 1 + docs/instrumentation.md | 1 + google/cloud/managed_spark_connect/session.py | 77 +++- tests/unit/test_session_start.py | 378 +++++++++++++++++- 4 files changed, 427 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index 04b769e2..f16e6775 100644 --- a/README.md +++ b/README.md @@ -139,6 +139,7 @@ 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. 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 ef5ec5f1..71b6e0c2 100644 --- a/docs/instrumentation.md +++ b/docs/instrumentation.md @@ -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/session.py b/google/cloud/managed_spark_connect/session.py index b45a3be3..6eb3fe27 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -277,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: @@ -774,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: @@ -848,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( @@ -977,6 +1007,21 @@ def create_session_pbar(): 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( @@ -1327,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 diff --git a/tests/unit/test_session_start.py b/tests/unit/test_session_start.py index 613c824b..9812c607 100644 --- a/tests/unit/test_session_start.py +++ b/tests/unit/test_session_start.py @@ -13,9 +13,11 @@ # 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 ServiceUnavailable +from google.api_core.exceptions import NotFound, ServiceUnavailable from google.cloud.dataproc_v1 import Session import pytest @@ -25,10 +27,13 @@ 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( @@ -43,6 +48,13 @@ def _failed(session_id, uuid, state=Session.State.FAILED, 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", [ @@ -52,15 +64,20 @@ def _failed(session_id, uuid, state=Session.State.FAILED, message=""): "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( - dataproc_mocks, state, message, detail + 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 = dataproc_mocks + mocks = no_retries mocks["client"].create_session.return_value.result.return_value = _failed( mocks["session_id"], mocks["uuid"], state, message ) @@ -71,18 +88,29 @@ def test_create_result_already_dead_raises_with_dataproc_state( 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 -def test_create_result_without_endpoint_waits_for_one(dataproc_mocks): +@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 = Session.State.CREATING - creating.runtime_info.endpoints = {} + 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] @@ -96,8 +124,8 @@ def test_create_result_without_endpoint_waits_for_one(dataproc_mocks): session.stop() -def test_waiting_for_endpoint_gives_up_when_the_session_fails(dataproc_mocks): - mocks = dataproc_mocks +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 = {} @@ -112,6 +140,52 @@ def test_waiting_for_endpoint_gives_up_when_the_session_fails(dataproc_mocks): 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 @@ -123,14 +197,12 @@ def first_rpc_fails(): yield spark_version -def test_unusable_new_session_fails_get_or_create( - dataproc_mocks, first_rpc_fails -): +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 = dataproc_mocks + mocks = no_retries mocks["client"].get_session.return_value = _failed( mocks["session_id"], mocks["uuid"], message="Agent failed to start" ) @@ -183,8 +255,8 @@ def test_new_session_is_returned_unless_dataproc_reports_it_gone( session.stop() -def test_new_session_check_is_recorded(dataproc_mocks, first_rpc_fails): - mocks = dataproc_mocks +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"] ) @@ -197,3 +269,281 @@ def test_new_session_check_is_recorded(dataproc_mocks, first_rpc_fails): 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"]