fix(spark): clean up session and port-forward on failure by shreya-024 · Pull Request #779 · kubeflow/sdk

shreya-024 · 2026-09-01T20:21:01Z

What this PR does / why we need it:

Fixes the SparkConnect resource handling gap identified in #476.

Map unrecognised CR states to SparkConnectState.UNKNOWN instead of PROVISIONING in get_spark_connect_info_from_cr
Delete the SparkConnect CR on failure in create_and_connect to prevent orphaned cluster resources
Terminate kubectl port-forward subprocess on connection/timeout failure
Added "unknown status" parametrized test case inside test_get_spark_connect_info_from_cr
Added test_create_and_connect_cleanup_on_failure in backend_test.py

Which issue(s) this PR fixes:

Fixes SparkConnect CR cleanup behavior on failure and SDK state mismatch #476

github-actions · 2026-09-01T20:21:42Z

🎉 Welcome to the Kubeflow SDK! 🎉

Thanks for opening your first PR! We're happy to have you as part of our community 🚀

Here's what happens next:

If you haven't already, please check out our Contributing Guide for repo-specific guidelines and the Kubeflow Contributor Guide for general community standards
Our team will review your PR soon! cc @kubeflow/kubeflow-sdk-team

Join the community:

Slack: Join our #kubeflow-ml-experience and #kubeflow-trainer Slack channels
Meetings: Attend the Kubeflow SDK and ML Experience bi-weekly meetings

Feel free to ask questions in the comments if you need any help or clarification!
Thanks again for contributing to Kubeflow! 🙏

tariq-hasan

Hi @shreya-024! I have added a few comments.

Please also follow these instructions to sign your commits: https://github.com/kubeflow/sdk/pull/779/checks?check_run_id=100012874640.

tariq-hasan · 2026-09-02T01:45:02Z

    RUNNING = "Running"  # Operator may set this when server is up; treated as ready
    NOT_READY = "NotReady"
    FAILED = "Failed"
+    UNKNOWN = "Unknown"


This should be removed.

Please remove UNKNOWN = "Unknown".

tariq-hasan · 2026-09-02T01:46:03Z

            state = SparkConnectState(spark_connect_cr.status.state)
        except ValueError:
-            state = SparkConnectState.PROVISIONING
+            state = SparkConnectState.UNKNOWN


Please use this instead.

sdk/kubeflow/common/constants.py

Lines 21 to 22 in 3febcd7

# Unknown indicates that the value can't be identified.

UNKNOWN = "Unknown"

tariq-hasan · 2026-09-02T01:57:52Z

The RUNNING state should be removed.

Suggested change

if info.state == SparkConnectState.READY:

logger.info(

tariq-hasan · 2026-09-02T01:59:12Z

@@ -32,6 +32,7 @@ class SparkConnectState(str, Enum):
    RUNNING = "Running"  # Operator may set this when server is up; treated as ready


The RUNNING state should be removed.

tariq-hasan · 2026-09-02T02:27:33Z

/ok-to-test
/retest

tariq-hasan · 2026-09-02T03:31:50Z

/assign

vjkumar2756

LGTM
@shreya-024
I went through the changes, and the cleanup handling looks solid. Mapping unrecognized states to UNKNOWN, cleaning up the SparkConnect CR on failure, and terminating the port forward subprocess should prevent orphaned resources and state mismatches.

tariq-hasan

Hi @shreya-024! I have added a few comments.

tariq-hasan · 2026-09-06T18:40:25Z

    RUNNING = "Running"  # Operator may set this when server is up; treated as ready
    NOT_READY = "NotReady"
    FAILED = "Failed"
+    UNKNOWN = "Unknown"


Please remove UNKNOWN = "Unknown".

tariq-hasan · 2026-09-06T19:01:02Z

Let's rewrite as follows.

Suggested change

state: str = common_constants.UNKNOWN

tariq-hasan · 2026-09-06T19:01:33Z

Same here

Suggested change

status: str = common_constants.UNKNOWN

tariq-hasan · 2026-09-06T19:02:08Z

            state = SparkConnectState(spark_connect_cr.status.state)
        except ValueError:
-            state = SparkConnectState.PROVISIONING
+            state = SparkConnectState(common_constants.UNKNOWN)


Suggested change

state = SparkConnectState(common_constants.UNKNOWN)

state = common_constants.UNKNOWN

tariq-hasan · 2026-09-06T19:59:45Z

We need to be able to retry for transient errors.

Please use the following implementation for _wait_for_session_ready. It handles transient errors with a retry mechanism.

def _wait_for_session_ready( self, name: str, timeout: int = 300, polling_interval: int = 2, ) -> SparkConnectInfo: """Wait for a SparkConnect session to become ready. Args: name: Name of the SparkConnect session. timeout: Maximum time in seconds to wait. polling_interval: Time in seconds between status checks. Returns: SparkConnectInfo containing information about the ready session. Raises: RuntimeError: If the SparkConnect session reaches the failed state. TimeoutError: If the session does not become ready within the timeout. Exception: If getting the session fails with a non-transient error. """ start_time = time.monotonic() last_log_time = start_time deadline = start_time + timeout while True: now = time.monotonic() if now >= deadline: raise TimeoutError( f"Timeout waiting for {constants.SPARK_CONNECT_KIND} to be ready: " f"{self.namespace}/{name} (timeout: {timeout}s)" ) try: info = self.get_session(name) except TimeoutError as e: # The request itself timed out. This is transient, so retry # as long as the overall wait deadline has not expired. logger.warning( "Timeout getting session %s/%s, retrying: %s", self.namespace, name, e, ) except RuntimeError as e: # get_session wraps ApiException in RuntimeError, so inspect the cause. cause = e.__cause__ if not isinstance(cause, client.ApiException) or cause.status not in { 429, 500, 502, 503, 504, }: raise logger.warning( "Transient HTTP %s getting session %s/%s, retrying: %s", cause.status, self.namespace, name, e, ) else: if info.state == SparkConnectState.READY: logger.info( "Session ready: %s/%s state=%s serviceName=%s (%.0fs)", self.namespace, name, info.state, info.service_name, time.monotonic() - start_time, ) return info if info.state == SparkConnectState.FAILED: raise RuntimeError( f"{constants.SPARK_CONNECT_KIND} failed: " f"{self.namespace}/{name}" ) now = time.monotonic() if now - last_log_time >= 10.0: logger.info( "Waiting for session: %s/%s state=%s serviceName=%s elapsed=%.0fs", self.namespace, name, info.state, info.service_name, now - start_time, ) last_log_time = now if now >= deadline: raise TimeoutError( f"Timeout waiting for {constants.SPARK_CONNECT_KIND} to be ready: " f"{self.namespace}/{name} (timeout: {timeout}s)" ) # Don't sleep beyond the overall deadline. remaining = deadline - time.monotonic() if remaining <= 0: raise TimeoutError( f"Timeout waiting for {constants.SPARK_CONNECT_KIND} to be ready: " f"{self.namespace}/{name} (timeout: {timeout}s)" ) time.sleep(min(polling_interval, remaining))

tariq-hasan

Hi @shreya-024! I have added a few more suggestions for changes.

tariq-hasan · 2026-09-11T01:08:54Z

+            )
+            with contextlib.suppress(Exception):
+                self.delete_session(info.name)
+            raise


We can rewrite this as follows. This is because Exception would not cover Ctrl+C keyboard interrupt for example. So to make the implementation more robust we use a finally block to always run session deletion if creation fails.

success = False try: info = self._wait_for_session_ready(info.name, timeout=timeout) logger.info( "Session ready, connecting (service_name=%s)", info.service_name, ) session = self.connect( info, connect_timeout=connect_timeout, ) success = True return session finally: if not success: logger.warning( "Failed to setup or connect to SparkConnect session %s/%s. " "Cleaning up SparkConnect session.", info.namespace, info.name, ) try: self.delete_session(info.name) except Exception: logger.exception( "Failed to clean up SparkConnect session %s/%s", info.namespace, info.name, )

tariq-hasan · 2026-09-12T14:22:20Z

Please use the following implementation for _wait_for_session_ready. It handles transient errors with a retry mechanism.

def _wait_for_session_ready( self, name: str, timeout: int = 300, polling_interval: int = 2, ) -> SparkConnectInfo: """Wait for a SparkConnect session to become ready. Args: name: Name of the SparkConnect session. timeout: Maximum time in seconds to wait. polling_interval: Time in seconds between status checks. Returns: SparkConnectInfo containing information about the ready session. Raises: RuntimeError: If the SparkConnect session reaches the failed state. TimeoutError: If the session does not become ready within the timeout. Exception: If getting the session fails with a non-transient error. """ start_time = time.monotonic() last_log_time = start_time deadline = start_time + timeout while True: now = time.monotonic() if now >= deadline: raise TimeoutError( f"Timeout waiting for {constants.SPARK_CONNECT_KIND} to be ready: " f"{self.namespace}/{name} (timeout: {timeout}s)" ) try: info = self.get_session(name) except TimeoutError as e: # The request itself timed out. This is transient, so retry # as long as the overall wait deadline has not expired. logger.warning( "Timeout getting session %s/%s, retrying: %s", self.namespace, name, e, ) except RuntimeError as e: # get_session wraps ApiException in RuntimeError, so inspect the cause. cause = e.__cause__ if not isinstance(cause, client.ApiException) or cause.status not in { 429, 500, 502, 503, 504, }: raise logger.warning( "Transient HTTP %s getting session %s/%s, retrying: %s", cause.status, self.namespace, name, e, ) else: if info.state == SparkConnectState.READY: logger.info( "Session ready: %s/%s state=%s serviceName=%s (%.0fs)", self.namespace, name, info.state, info.service_name, time.monotonic() - start_time, ) return info if info.state == SparkConnectState.FAILED: raise RuntimeError( f"{constants.SPARK_CONNECT_KIND} failed: " f"{self.namespace}/{name}" ) now = time.monotonic() if now - last_log_time >= 10.0: logger.info( "Waiting for session: %s/%s state=%s serviceName=%s elapsed=%.0fs", self.namespace, name, info.state, info.service_name, now - start_time, ) last_log_time = now if now >= deadline: raise TimeoutError( f"Timeout waiting for {constants.SPARK_CONNECT_KIND} to be ready: " f"{self.namespace}/{name} (timeout: {timeout}s)" ) # Don't sleep beyond the overall deadline. remaining = deadline - time.monotonic() if remaining <= 0: raise TimeoutError( f"Timeout waiting for {constants.SPARK_CONNECT_KIND} to be ready: " f"{self.namespace}/{name} (timeout: {timeout}s)" ) time.sleep(min(polling_interval, remaining))

tariq-hasan · 2026-09-12T14:42:08Z

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


Unit tests need to use the existing test_<method> function naming pattern with test cases using the TestCase helper.

Please also add test cases that cover

create succeeds, connect succeeds

create succeeds, wait fails

create succeeds, connect fails

create itself fails

Please move all these test cases into test_create_and_connect.

The pattern we follow is to group test cases under test_<method_name> in each unit test module.

tariq-hasan · 2026-09-12T15:11:53Z


    # Parse state
-    state = SparkConnectState.PROVISIONING
+    state = common_constants.UNKNOWN


Actually I think it's better for this one to be state = SparkConnectState.PROVISIONING.

Granted the PROVISIONING state is not currently leveraged in the operator. At the same time it makes sense for the initial state on the client side to be a provisioning state and not an unknown state.

Please also remove the ready state from the following since only the provisioning state can be captured after CR creation before it is read back out from the control plane.

sdk/test/e2e/spark/test_spark_examples.py

Line 114 in 3febcd7

assert info.state in (SparkConnectState.PROVISIONING, SparkConnectState.READY)

shreya-024 · 2026-09-16T10:31:44Z

Hi @tariq-hasan, I've addressed all the review comments. Could you please take a look when you get a chance?

tariq-hasan · 2026-10-04T19:53:37Z

Thanks @shreya-024! Please sign the DCO commit as well: https://github.com/kubeflow/sdk/pull/779/checks?check_run_id=111507077564.

Signed-off-by: shreya-024 <shreyak.2406@gmail.com>

shreya-024 · 2026-10-04T21:38:51Z

@tariq-hasan I've signed the commit

tariq-hasan · 2026-10-07T15:03:17Z

Thanks!
/lgtm

andreyvelich

/lgtm
/approve

google-oss-prow · 2026-10-07T16:06:01Z

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: andreyvelich

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Details

Needs approval from an approver in each of these files:

~~OWNERS~~ [andreyvelich]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

google-oss-prow Bot requested review from Fiona-Waters and tariq-hasan September 1, 2026 20:21

google-oss-prow Bot added the size/L label Sep 1, 2026

shreya-024 mentioned this pull request Sep 1, 2026

SparkConnect CR cleanup behavior on failure and SDK state mismatch #476

Closed

tariq-hasan reviewed Sep 2, 2026

View reviewed changes

google-oss-prow Bot added the ok-to-test label Sep 2, 2026

google-oss-prow Bot assigned tariq-hasan Sep 2, 2026

shreya-024 force-pushed the fix/spark-connect-state-unknown-cleanup branch from 8d74bdc to 92e552c Compare September 2, 2026 05:40

vjkumar2756 reviewed Sep 2, 2026

View reviewed changes

shreya-024 requested a review from tariq-hasan September 4, 2026 20:05

tariq-hasan suggested changes Sep 6, 2026

View reviewed changes

tariq-hasan suggested changes Sep 12, 2026

View reviewed changes

shreya-024 force-pushed the fix/spark-connect-state-unknown-cleanup branch from fdb5397 to 5fb36b3 Compare September 16, 2026 10:28

shreya-024 requested a review from tariq-hasan September 16, 2026 10:32

RobuRishabh added this to Kubeflow SDK Sep 21, 2026

RobuRishabh moved this to Needs Triage in Kubeflow SDK Sep 21, 2026

tariq-hasan moved this from Needs Triage to In Progress in Kubeflow SDK Sep 21, 2026

shreya-024 added 4 commits October 5, 2026 01:51

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

3a0427c

Signed-off-by: shreya-024 <shreyak.2406@gmail.com>

fix(spark): remove RUNNING state and use constant.UNKNOWN fallback

8108979

Signed-off-by: shreya-024 <shreyak.2406@gmail.com>

fix(spark): address review comments on state and constants

3a84d00

Signed-off-by: shreya-024 <shreyak.2406@gmail.com>

test(spark): group create and connect cases

b96c93d

Signed-off-by: shreya-024 <shreyak.2406@gmail.com>

shreya-024 force-pushed the fix/spark-connect-state-unknown-cleanup branch from 8acdf84 to b96c93d Compare October 4, 2026 20:21

google-oss-prow Bot added the lgtm label Oct 7, 2026

andreyvelich reviewed Oct 7, 2026

View reviewed changes

google-oss-prow Bot assigned andreyvelich Oct 7, 2026

google-oss-prow Bot added the approved label Oct 7, 2026

google-oss-prow Bot merged commit 88f7395 into kubeflow:main Oct 7, 2026
19 checks passed

google-oss-prow Bot added this to the v0.6 milestone Oct 7, 2026

tariq-hasan mentioned this pull request Oct 7, 2026

fix(spark): clean up SparkConnect CR on failure and align states #676

Closed

3 tasks done

	# Unknown indicates that the value can't be identified.
	UNKNOWN = "Unknown"

		@@ -32,6 +32,7 @@ class SparkConnectState(str, Enum):
		RUNNING = "Running" # Operator may set this when server is up; treated as ready

	state = SparkConnectState(common_constants.UNKNOWN)
	state = common_constants.UNKNOWN

Conversation

shreya-024 commented Sep 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Sep 1, 2026

Uh oh!

tariq-hasan left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tariq-hasan commented Sep 2, 2026

Uh oh!

tariq-hasan commented Sep 2, 2026

Uh oh!

vjkumar2756 left a comment

Choose a reason for hiding this comment

Uh oh!

tariq-hasan left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tariq-hasan left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shreya-024 commented Sep 16, 2026

Uh oh!

tariq-hasan commented Oct 4, 2026

Uh oh!

shreya-024 commented Oct 4, 2026

Uh oh!

tariq-hasan commented Oct 7, 2026

Uh oh!

andreyvelich left a comment

Choose a reason for hiding this comment

Uh oh!

google-oss-prow Bot commented Oct 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

shreya-024 commented Sep 1, 2026 •

edited

Loading

tariq-hasan left a comment •

edited

Loading