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

kubeflow/spark/backends/kubernetes/backend.py

@@ -336,7 +336,7 @@ def _wait_for_session_ready(
         while True:
             info = self.get_session(name)
 
-            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 +601,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 +685,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

@@ -632,7 +632,7 @@ def get_spark_connect_info_from_cr(
         try:
             state = SparkConnectState(spark_connect_cr.status.state)
         except ValueError:
-            state = SparkConnectState.PROVISIONING
+            state = SparkConnectState(common_constants.UNKNOWN)
 
     # Extract server status
     server_status = None

kubeflow/spark/backends/kubernetes/utils_test.py

@@ -820,6 +820,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,
@@ -870,13 +883,16 @@ def test_get_spark_connect_info_from_cr(
             assert info.state == SparkConnectState.FAILED
 
         elif test_case.name == "running status":
-            assert info.state == SparkConnectState.RUNNING
+            assert info.state == SparkConnectState.UNKNOWN
             assert info.service_name == "run-session-svc"
 
         elif test_case.name == "empty status":
             assert info.state == SparkConnectState.PROVISIONING
             assert info.driver_pod_name is None
 
+        elif test_case.name == "unknown status":
+            assert info.state == SparkConnectState.UNKNOWN
+
     else:
         with pytest.raises(
             test_case.expected_error,

kubeflow/spark/types/types.py

@@ -29,9 +29,9 @@ 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"
+    UNKNOWN = "Unknown"
 
 
 @dataclass

kubeflow/spark/types/types_test.py

@@ -37,7 +37,6 @@
     [
         (SparkConnectState.PROVISIONING, "Provisioning"),
         (SparkConnectState.READY, "Ready"),
-        (SparkConnectState.RUNNING, "Running"),
         (SparkConnectState.NOT_READY, "NotReady"),
         (SparkConnectState.FAILED, "Failed"),
     ],
@@ -52,7 +51,6 @@ def test_spark_connect_state_values(state, expected):
     [
         SparkConnectState.PROVISIONING,
         SparkConnectState.READY,
-        SparkConnectState.RUNNING,
         SparkConnectState.NOT_READY,
         SparkConnectState.FAILED,
     ],