You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by do...@apache.org on 2023/09/17 09:02:14 UTC
[spark] branch master updated: [SPARK-45167][CONNECT][PYTHON] Python client must call `release_all`
This is an automated email from the ASF dual-hosted git repository.
dongjoon pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 65b12f95c14 [SPARK-45167][CONNECT][PYTHON] Python client must call `release_all`
65b12f95c14 is described below
commit 65b12f95c14935528d736d9f7c8e8e181a7f16f8
Author: Martin Grund <ma...@databricks.com>
AuthorDate: Sun Sep 17 02:02:05 2023 -0700
[SPARK-45167][CONNECT][PYTHON] Python client must call `release_all`
### What changes were proposed in this pull request?
Previously the Python client would not call `release_all` after fetching all results and leaving the query dangling. The query would then be removed after the five minute timeout.
This patch adds proper testing for calling release all and release until.
In addition it fixes a test race condition where we would close the SparkSession which would in turn close the GRPC channel which might have dangling async release calls hanging.
### Why are the changes needed?
Stabililty
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
New UT
### Was this patch authored or co-authored using generative AI tooling?
No
Closes #42929 from grundprinzip/SPARK-45167.
Lead-authored-by: Martin Grund <ma...@databricks.com>
Co-authored-by: Martin Grund <gr...@gmail.com>
Signed-off-by: Dongjoon Hyun <dh...@apple.com>
---
python/pyspark/sql/connect/client/core.py | 1 +
python/pyspark/sql/connect/client/reattach.py | 37 ++++++++++--
.../sql/tests/connect/client/test_client.py | 70 +++++++++++++++++-----
3 files changed, 87 insertions(+), 21 deletions(-)
diff --git a/python/pyspark/sql/connect/client/core.py b/python/pyspark/sql/connect/client/core.py
index 0584f18e5bb..9bf3eb3f0ac 100644
--- a/python/pyspark/sql/connect/client/core.py
+++ b/python/pyspark/sql/connect/client/core.py
@@ -1017,6 +1017,7 @@ class SparkConnectClient(object):
"""
Close the channel.
"""
+ ExecutePlanResponseReattachableIterator.shutdown()
self._channel.close()
self._closed = True
diff --git a/python/pyspark/sql/connect/client/reattach.py b/python/pyspark/sql/connect/client/reattach.py
index 7e1e722d5fd..e58864b965b 100644
--- a/python/pyspark/sql/connect/client/reattach.py
+++ b/python/pyspark/sql/connect/client/reattach.py
@@ -21,7 +21,9 @@ check_dependencies(__name__)
import warnings
import uuid
from collections.abc import Generator
-from typing import Optional, Dict, Any, Iterator, Iterable, Tuple, Callable, cast
+from typing import Optional, Dict, Any, Iterator, Iterable, Tuple, Callable, cast, Type, ClassVar
+from multiprocessing import RLock
+from multiprocessing.synchronize import RLock as RLockBase
from multiprocessing.pool import ThreadPool
import os
@@ -53,7 +55,30 @@ class ExecutePlanResponseReattachableIterator(Generator):
ReleaseExecute RPCs that instruct the server to release responses that it already processed.
"""
- _release_thread_pool = ThreadPool(os.cpu_count() if os.cpu_count() else 8)
+ # Lock to manage the pool
+ _lock: ClassVar[RLockBase] = RLock()
+ _release_thread_pool: Optional[ThreadPool] = ThreadPool(os.cpu_count() if os.cpu_count() else 8)
+
+ @classmethod
+ def shutdown(cls: Type["ExecutePlanResponseReattachableIterator"]) -> None:
+ """
+ When the channel is closed, this method will be called before, to make sure all
+ outstanding calls are closed.
+ """
+ with cls._lock:
+ if cls._release_thread_pool is not None:
+ cls._release_thread_pool.close()
+ cls._release_thread_pool.join()
+ cls._release_thread_pool = None
+
+ @classmethod
+ def _initialize_pool_if_necessary(cls: Type["ExecutePlanResponseReattachableIterator"]) -> None:
+ """
+ If the processing pool for the release calls is None, initialize the pool exactly once.
+ """
+ with cls._lock:
+ if cls._release_thread_pool is None:
+ cls._release_thread_pool = ThreadPool(os.cpu_count() if os.cpu_count() else 8)
def __init__(
self,
@@ -62,6 +87,7 @@ class ExecutePlanResponseReattachableIterator(Generator):
retry_policy: Dict[str, Any],
metadata: Iterable[Tuple[str, str]],
):
+ ExecutePlanResponseReattachableIterator._initialize_pool_if_necessary()
self._request = request
self._retry_policy = retry_policy
if request.operation_id:
@@ -111,7 +137,6 @@ class ExecutePlanResponseReattachableIterator(Generator):
self._last_returned_response_id = ret.response_id
if ret.HasField("result_complete"):
- self._result_complete = True
self._release_all()
else:
self._release_until(self._last_returned_response_id)
@@ -190,7 +215,8 @@ class ExecutePlanResponseReattachableIterator(Generator):
except Exception as e:
warnings.warn(f"ReleaseExecute failed with exception: {e}.")
- ExecutePlanResponseReattachableIterator._release_thread_pool.apply_async(target)
+ if ExecutePlanResponseReattachableIterator._release_thread_pool is not None:
+ ExecutePlanResponseReattachableIterator._release_thread_pool.apply_async(target)
def _release_all(self) -> None:
"""
@@ -218,7 +244,8 @@ class ExecutePlanResponseReattachableIterator(Generator):
except Exception as e:
warnings.warn(f"ReleaseExecute failed with exception: {e}.")
- ExecutePlanResponseReattachableIterator._release_thread_pool.apply_async(target)
+ if ExecutePlanResponseReattachableIterator._release_thread_pool is not None:
+ ExecutePlanResponseReattachableIterator._release_thread_pool.apply_async(target)
self._result_complete = True
def _call_iter(self, iter_fun: Callable) -> Any:
diff --git a/python/pyspark/sql/tests/connect/client/test_client.py b/python/pyspark/sql/tests/connect/client/test_client.py
index 70280c1d24a..5fe4f7f4637 100644
--- a/python/pyspark/sql/tests/connect/client/test_client.py
+++ b/python/pyspark/sql/tests/connect/client/test_client.py
@@ -136,9 +136,12 @@ class SparkConnectClientReattachTestCase(unittest.TestCase):
"jitter": 10,
"min_jitter_threshold": 10,
}
- self.response = proto.ExecutePlanResponse()
+ self.response = proto.ExecutePlanResponse(
+ response_id="1",
+ )
self.finished = proto.ExecutePlanResponse(
- result_complete=proto.ExecutePlanResponse.ResultComplete()
+ result_complete=proto.ExecutePlanResponse.ResultComplete(),
+ response_id="2",
)
def _stub_with(self, execute=None, attach=None):
@@ -147,15 +150,33 @@ class SparkConnectClientReattachTestCase(unittest.TestCase):
attach_ops=ResponseGenerator(attach) if attach is not None else None,
)
+ def assertEventually(self, callable, timeout_ms=1000):
+ """Helper method that will continuously evaluate the callable to not raise an
+ exception."""
+ import time
+
+ limit = time.monotonic_ns() + timeout_ms * 1000 * 1000
+ while time.monotonic_ns() < limit:
+ try:
+ callable()
+ break
+ except Exception:
+ time.sleep(0.1)
+ callable()
+
def test_basic_flow(self):
stub = self._stub_with([self.response, self.finished])
ite = ExecutePlanResponseReattachableIterator(self.request, stub, self.policy, [])
for b in ite:
pass
- self.assertEqual(0, stub.attach_calls)
- self.assertGreater(1, stub.release_calls)
- self.assertEqual(1, stub.execute_calls)
+ def check_all():
+ self.assertEqual(0, stub.attach_calls)
+ self.assertEqual(1, stub.release_until_calls)
+ self.assertEqual(1, stub.release_calls)
+ self.assertEqual(1, stub.execute_calls)
+
+ self.assertEventually(check_all, timeout_ms=1000)
def test_fail_during_execute(self):
def fatal():
@@ -167,9 +188,13 @@ class SparkConnectClientReattachTestCase(unittest.TestCase):
for b in ite:
pass
- self.assertEqual(0, stub.attach_calls)
- self.assertEqual(0, stub.release_calls)
- self.assertEqual(1, stub.execute_calls)
+ def check():
+ self.assertEqual(0, stub.attach_calls)
+ self.assertEqual(1, stub.release_calls)
+ self.assertEqual(1, stub.release_until_calls)
+ self.assertEqual(1, stub.execute_calls)
+
+ self.assertEventually(check, timeout_ms=1000)
def test_fail_and_retry_during_execute(self):
def non_fatal():
@@ -182,9 +207,13 @@ class SparkConnectClientReattachTestCase(unittest.TestCase):
for b in ite:
pass
- self.assertEqual(1, stub.attach_calls)
- self.assertEqual(1, stub.release_calls)
- self.assertEqual(1, stub.execute_calls)
+ def check():
+ self.assertEqual(1, stub.attach_calls)
+ self.assertEqual(1, stub.release_calls)
+ self.assertEqual(3, stub.release_until_calls)
+ self.assertEqual(1, stub.execute_calls)
+
+ self.assertEventually(check, timeout_ms=1000)
def test_fail_and_retry_during_reattach(self):
count = 0
@@ -204,9 +233,13 @@ class SparkConnectClientReattachTestCase(unittest.TestCase):
for b in ite:
pass
- self.assertEqual(2, stub.attach_calls)
- self.assertEqual(2, stub.release_calls)
- self.assertEqual(1, stub.execute_calls)
+ def check():
+ self.assertEqual(2, stub.attach_calls)
+ self.assertEqual(3, stub.release_until_calls)
+ self.assertEqual(1, stub.release_calls)
+ self.assertEqual(1, stub.execute_calls)
+
+ self.assertEventually(check, timeout_ms=1000)
class TestException(grpc.RpcError, grpc.Call):
@@ -257,6 +290,7 @@ class MockSparkConnectStub:
# Call counters
self.execute_calls = 0
self.release_calls = 0
+ self.release_until_calls = 0
self.attach_calls = 0
def ExecutePlan(self, *args, **kwargs):
@@ -267,8 +301,12 @@ class MockSparkConnectStub:
self.attach_calls += 1
return self._attach_ops
- def ReleaseExecute(self, *args, **kwargs):
- self.release_calls += 1
+ def ReleaseExecute(self, req: proto.ReleaseExecuteRequest, *args, **kwargs):
+ if req.HasField("release_all"):
+ self.release_calls += 1
+ elif req.HasField("release_until"):
+ print("increment")
+ self.release_until_calls += 1
class MockService:
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org