diff --git a/kubeflow/spark/backends/kubernetes/backend.py b/kubeflow/spark/backends/kubernetes/backend.py index 6d75874c8..e34d9d3dc 100644 --- a/kubeflow/spark/backends/kubernetes/backend.py +++ b/kubeflow/spark/backends/kubernetes/backend.py @@ -334,9 +334,24 @@ def _wait_for_session_ready( last_log_time = start_time while True: - info = self.get_session(name) + try: + info = self.get_session(name) + except Exception as e: + logger.warning( + "Transient error getting session %s/%s, retrying: %s", + self.namespace, + name, + e, + ) + if time.monotonic() - start_time >= timeout: + raise TimeoutError( + f"Timeout waiting for {constants.SPARK_CONNECT_KIND} to be ready: " + f"{self.namespace}/{name} (timeout: {timeout}s)" + ) from e + time.sleep(polling_interval) + continue - if info.state in (SparkConnectState.READY, SparkConnectState.RUNNING): + if info.state == SparkConnectState.READY: logger.info( "Session ready: %s/%s state=%s serviceName=%s (%.0fs)", self.namespace, @@ -601,29 +616,36 @@ def _get_or_create() -> None: thread.start() thread.join(timeout=connect_timeout) - if not thread.is_alive(): - if exc_holder: - raise exc_holder[0] - if result: - return result[0] - - # Connection timed out - base_msg = ( - f"Spark Connect connection to {connect_url} did not complete " - f"within {connect_timeout}s. " - "Verify: (1) port-forward target is the Spark Connect server pod, " - "(2) PySpark and server Spark major.minor match, " - "(3) driver pod logs for gRPC/auth errors; " - "see Spark sql/connect for server config." - ) - if pf_proc is not None and pf_proc.poll() is not None: - stderr_b = pf_proc.stderr.read() if pf_proc.stderr else b"" - stderr_str = stderr_b.decode("utf-8", errors="replace").strip() if stderr_b else "" - base_msg += ( - f" Port-forward process exited during connect " - f"(code={pf_proc.returncode}). stderr: {stderr_str}" + try: + if not thread.is_alive(): + if exc_holder: + raise exc_holder[0] + if result: + return result[0] + + # Connection timed out + base_msg = ( + f"Spark Connect connection to {connect_url} did not complete " + f"within {connect_timeout}s. " + "Verify: (1) port-forward target is the Spark Connect server pod, " + "(2) PySpark and server Spark major.minor match, " + "(3) driver pod logs for gRPC/auth errors; " + "see Spark sql/connect for server config." ) - raise TimeoutError(base_msg) + if pf_proc is not None and pf_proc.poll() is not None: + stderr_b = pf_proc.stderr.read() if pf_proc.stderr else b"" + stderr_str = stderr_b.decode("utf-8", errors="replace").strip() if stderr_b else "" + base_msg += ( + f" Port-forward process exited during connect " + f"(code={pf_proc.returncode}). stderr: {stderr_str}" + ) + raise TimeoutError(base_msg) + except Exception: + if pf_proc is not None and pf_proc.poll() is None: + pf_proc.terminate() + with contextlib.suppress(Exception): + pf_proc.wait(timeout=2) + raise def create_and_connect( self, @@ -678,10 +700,21 @@ def create_and_connect( timeout, ) - info = self._wait_for_session_ready(info.name, timeout=timeout) - logger.info("Session ready, connecting (service_name=%s)", info.service_name) - - return self.connect(info, connect_timeout=connect_timeout) + try: + info = self._wait_for_session_ready(info.name, timeout=timeout) + logger.info("Session ready, connecting (service_name=%s)", info.service_name) + return self.connect(info, connect_timeout=connect_timeout) + except Exception as e: + logger.warning( + "Failed to setup or connect to SparkConnect session %s/%s: %s. " + "Cleaning up SparkConnect session.", + info.namespace, + info.name, + e, + ) + with contextlib.suppress(Exception): + self.delete_session(info.name) + raise def get_session_logs( self, diff --git a/kubeflow/spark/backends/kubernetes/backend_test.py b/kubeflow/spark/backends/kubernetes/backend_test.py index 04d422de0..d2e61326a 100644 --- a/kubeflow/spark/backends/kubernetes/backend_test.py +++ b/kubeflow/spark/backends/kubernetes/backend_test.py @@ -16,7 +16,7 @@ from datetime import datetime import multiprocessing -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch from kubeflow_spark_api import models from kubernetes import client @@ -1941,3 +1941,24 @@ def test_get_job_logs(kubernetes_backend, test_case): raise print("test execution complete") + + +def test_create_and_connect_cleanup_on_failure(kubernetes_backend): + """Test that create_and_connect cleans up session when connect fails.""" + mock_info = MagicMock() + mock_info.name = "test-session" + mock_info.namespace = "default" + + with ( + patch.object(kubernetes_backend, "_create_session", return_value=mock_info), + patch.object( + kubernetes_backend, + "_wait_for_session_ready", + side_effect=RuntimeError("Wait failed"), + ), + patch.object(kubernetes_backend, "delete_session") as mock_delete, + ): + with pytest.raises(RuntimeError, match="Wait failed"): + kubernetes_backend.create_and_connect() + + mock_delete.assert_called_once_with("test-session") diff --git a/kubeflow/spark/backends/kubernetes/utils.py b/kubeflow/spark/backends/kubernetes/utils.py index 1885bce04..50b8c7bb8 100644 --- a/kubeflow/spark/backends/kubernetes/utils.py +++ b/kubeflow/spark/backends/kubernetes/utils.py @@ -627,12 +627,12 @@ def get_spark_connect_info_from_cr( raise ValueError(f"SparkConnect CR is invalid: {spark_connect_cr}") # Parse state - state = SparkConnectState.PROVISIONING + state = common_constants.UNKNOWN if spark_connect_cr.status and spark_connect_cr.status.state: try: state = SparkConnectState(spark_connect_cr.status.state) except ValueError: - state = SparkConnectState.PROVISIONING + state = common_constants.UNKNOWN # Extract server status server_status = None diff --git a/kubeflow/spark/backends/kubernetes/utils_test.py b/kubeflow/spark/backends/kubernetes/utils_test.py index 35170d3b7..e29e3d587 100644 --- a/kubeflow/spark/backends/kubernetes/utils_test.py +++ b/kubeflow/spark/backends/kubernetes/utils_test.py @@ -21,6 +21,7 @@ from kubeflow_spark_api import models import pytest +from kubeflow.common import constants as common_constants from kubeflow.spark.backends.kubernetes import constants from kubeflow.spark.backends.kubernetes.backend import KubernetesBackend from kubeflow.spark.backends.kubernetes.utils import ( @@ -820,6 +821,19 @@ def test_build_spark_connect_cr(test_case: TestCase, mock_k8s_backend) -> None: }, }, ), + TestCase( + name="unknown status", + expected_status=SUCCESS, + config={ + "metadata": { + "name": "unknown-session", + "namespace": "default", + }, + "status": models.SparkV1alpha1SparkConnectStatus( + state="InvalidOrUnknownState", + ), + }, + ), TestCase( name="missing name", expected_status=FAILED, @@ -861,21 +875,23 @@ def test_get_spark_connect_info_from_cr( assert info.service_name == "my-session-svc" assert info.creation_timestamp is not None - elif test_case.name == "provisioning status": - assert info.name == "new-session" - assert info.namespace == "spark" - assert info.state == SparkConnectState.PROVISIONING + elif test_case.name == "empty status": + assert info.state == common_constants.UNKNOWN # not PROVISIONING + assert info.driver_pod_name is None elif test_case.name == "failed status": assert info.state == SparkConnectState.FAILED elif test_case.name == "running status": - assert info.state == SparkConnectState.RUNNING + assert info.state == common_constants.UNKNOWN assert info.service_name == "run-session-svc" - - elif test_case.name == "empty status": + elif test_case.name == "provisioning status": + assert info.name == "new-session" + assert info.namespace == "spark" assert info.state == SparkConnectState.PROVISIONING - assert info.driver_pod_name is None + + elif test_case.name == "unknown status": + assert info.state == common_constants.UNKNOWN else: with pytest.raises( diff --git a/kubeflow/spark/types/types.py b/kubeflow/spark/types/types.py index e1779139d..4e617b848 100644 --- a/kubeflow/spark/types/types.py +++ b/kubeflow/spark/types/types.py @@ -21,6 +21,8 @@ import logging from typing import Any +import kubeflow.common.constants as common_constants + logger = logging.getLogger(__name__) @@ -29,7 +31,6 @@ class SparkConnectState(str, Enum): PROVISIONING = "Provisioning" READY = "Ready" - RUNNING = "Running" # Operator may set this when server is up; treated as ready NOT_READY = "NotReady" FAILED = "Failed" @@ -51,7 +52,7 @@ class SparkConnectInfo: name: str namespace: str - state: SparkConnectState + state: str = common_constants.UNKNOWN driver_pod_name: str | None = None pod_ip: str | None = None service_name: str | None = None @@ -206,7 +207,7 @@ class SparkJob: name: str namespace: str - status: SparkJobStatus | None = None + status: str = common_constants.UNKNOWN creation_timestamp: datetime | None = None num_executors: int | None = None driver_pod_name: str | None = None diff --git a/kubeflow/spark/types/types_test.py b/kubeflow/spark/types/types_test.py index f32929943..66a8f3305 100644 --- a/kubeflow/spark/types/types_test.py +++ b/kubeflow/spark/types/types_test.py @@ -19,6 +19,7 @@ import pytest +from kubeflow.common import constants as common_constants from kubeflow.spark.types.types import ( Driver, Executor, @@ -37,7 +38,6 @@ [ (SparkConnectState.PROVISIONING, "Provisioning"), (SparkConnectState.READY, "Ready"), - (SparkConnectState.RUNNING, "Running"), (SparkConnectState.NOT_READY, "NotReady"), (SparkConnectState.FAILED, "Failed"), ], @@ -52,7 +52,6 @@ def test_spark_connect_state_values(state, expected): [ SparkConnectState.PROVISIONING, SparkConnectState.READY, - SparkConnectState.RUNNING, SparkConnectState.NOT_READY, SparkConnectState.FAILED, ], @@ -355,7 +354,7 @@ def test_spark_job(test_case: TestCase): assert getattr(job, key) == value if test_case.name == "default spark job": - assert job.status is None + assert job.status == common_constants.UNKNOWN assert job.creation_timestamp is None assert job.num_executors is None assert job.driver_pod_name is None