Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
google-oss-prow merged 4 commits into from
Oct 7, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3a0427c
fix(spark): clean up session and port-forward on failure
shreya-024 Sep 1, 2026
8108979
fix(spark): remove RUNNING state and use constant.UNKNOWN fallback
shreya-024 Sep 2, 2026
3a84d00
fix(spark): address review comments on state and constants
shreya-024 Sep 7, 2026
b96c93d
test(spark): group create and connect cases
shreya-024 Oct 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
197 changes: 138 additions & 59 deletions kubeflow/spark/backends/kubernetes/backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -329,66 +329,107 @@ def _wait_for_session_ready(

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:
info = self.get_session(name)
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)"
)

if info.state in (SparkConnectState.READY, SparkConnectState.RUNNING):
logger.info(
"Session ready: %s/%s state=%s serviceName=%s (%.0fs)",
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,
info.state,
info.service_name,
time.monotonic() - start_time,
)
return info

if info.state == SparkConnectState.FAILED:
raise RuntimeError(
f"{constants.SPARK_CONNECT_KIND} failed: {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

now = time.monotonic()
if now - last_log_time >= 10.0:
logger.info(
"Waiting for session: %s/%s state=%s serviceName=%s elapsed=%.0fs",
logger.warning(
"Transient HTTP %s getting session %s/%s, retrying: %s",
cause.status,
self.namespace,
name,
info.state,
info.service_name,
now - start_time,
e,
)
last_log_time = now
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 now - start_time >= timeout:
if info.state == SparkConnectState.FAILED:
raise RuntimeError(
f"{constants.SPARK_CONNECT_KIND} failed: {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(polling_interval)
time.sleep(min(polling_interval, remaining))

def _wait_for_connect_port(
self, host: str, port: int, timeout_sec: int = 60, interval_sec: float = 2.0
) -> bool:
"""Wait until a Spark Connect server becomes reachable.

Args:
host:
Hostname or IP address of the Spark Connect server.

port:
TCP port of the Spark Connect server.

timeout_sec:
Maximum time in seconds to wait.

interval_sec:
Time in seconds between connection attempts.
host: Hostname or IP address of the Spark Connect server.
port: TCP port of the Spark Connect server.
timeout_sec: Maximum time in seconds to wait.
interval_sec: Time in seconds between connection attempts.

Returns:
True if the server becomes reachable before the timeout, otherwise False.
Expand Down Expand Up @@ -601,29 +642,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,
Expand Down Expand Up @@ -678,10 +726,41 @@ 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)
success = False

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

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

def get_session_logs(
self,
Expand Down
91 changes: 65 additions & 26 deletions kubeflow/spark/backends/kubernetes/backend_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -852,40 +852,79 @@ def test_validate_spark_connect_url(test_case):
expected_status=SUCCESS,
config={"session_name": None},
),
TestCase(
name="create succeeds, wait fails",
expected_status=FAILED,
config={"session_name": "test-session", "wait_error": RuntimeError("Wait failed")},
expected_error=RuntimeError,
),
TestCase(
name="create succeeds, connect fails",
expected_status=FAILED,
config={
"session_name": "test-session",
"connect_error": RuntimeError("Connect failed"),
},
expected_error=RuntimeError,
),
TestCase(
name="create itself fails",
expected_status=FAILED,
config={"session_name": "test-session", "create_error": RuntimeError("Create failed")},
expected_error=RuntimeError,
),
],
)
def test_create_and_connect(kubernetes_backend, test_case):
"""Test create_and_connect with and without Name option."""
print("Executing test:", test_case.name)
try:
options = (
[Name(test_case.config["session_name"])] if test_case.config["session_name"] else None
)
ready_info = SparkConnectInfo(
name=test_case.config["session_name"] or "spark-connect-abc",
namespace=DEFAULT_NAMESPACE,
state=SparkConnectState.READY,
service_name="svc",
)
options = [Name(test_case.config["session_name"])] if test_case.config["session_name"] else None
session_name = test_case.config["session_name"] or "spark-connect-abc"
ready_info = SparkConnectInfo(
name=session_name,
namespace=DEFAULT_NAMESPACE,
state=SparkConnectState.READY,
service_name="svc",
)

with (
patch.object(
kubernetes_backend, "_create_session", return_value=ready_info
) as mock_create,
patch.object(kubernetes_backend, "_wait_for_session_ready", return_value=ready_info),
patch.object(
kubernetes_backend, "get_connect_url", return_value=("sc://localhost:15002", None)
),
patch("kubeflow.spark.backends.kubernetes.backend.SparkSession"),
):
create_error = test_case.config.get("create_error")
wait_error = test_case.config.get("wait_error")
connect_error = test_case.config.get("connect_error")
with (
patch.object(
kubernetes_backend,
"_create_session",
side_effect=create_error,
return_value=ready_info,
) as mock_create,
patch.object(
kubernetes_backend,
"_wait_for_session_ready",
side_effect=wait_error,
return_value=ready_info,
),
patch.object(
kubernetes_backend,
"connect",
side_effect=connect_error,
return_value=MagicMock(),
),
patch.object(kubernetes_backend, "delete_session") as mock_delete,
):
try:
kubernetes_backend.create_and_connect(options=options)
assert test_case.expected_status == SUCCESS
mock_delete.assert_not_called()
except Exception as e:
assert type(e) is test_case.expected_error
if create_error:
mock_delete.assert_not_called()
else:
mock_delete.assert_called_once_with(session_name)

if not create_error:
mock_create.assert_called_once()
assert mock_create.call_args.kwargs.get("options") == options

assert test_case.expected_status == SUCCESS

except Exception as e:
assert type(e) is test_case.expected_error
print("test execution complete")


Expand Down
2 changes: 1 addition & 1 deletion kubeflow/spark/backends/kubernetes/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = common_constants.UNKNOWN

# Extract server status
server_status = None
Expand Down
Loading
Loading