fix(spark): clean up session and port-forward on failure

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,

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")

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

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(

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

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