dongjoon-hyun commented on code in PR #58396:
URL: https://github.com/apache/spark/pull/58396#discussion_r4420284408
##########
python/pyspark/sql/connect/client/core.py:
##########
@@ -2604,14 +2606,13 @@ def _delete_ml_cache(self, cache_ids: List[str],
evict_only: bool = False) -> Li
return []
def _on_exit(self) -> None:
- # If the client has already been explicitly closed, skip all cleanup
RPCs.
- # The server-side resources were released by close(); reissuing them
here
- # is wasted work and, if the server has since become unreachable, can
- # block process exit on the gRPC call.
- if self._closed:
+ # Skip cleanup RPCs if the session is already closed or released.
+ if self._closed or self._session_released:
return
- self._cleanup_ml_cache()
+ # No server-side session ID (None or "") means no ML cache to clean.
+ if self._server_session_id:
Review Comment:
**2. [Low] An empty server-side session ID could be dropped where it is
stored**
This check treats `""` like `None`, but the other readers of
`_server_session_id` use `is not None`, for example
`_execute_plan_request_with_metadata`
([L1547-L1548](https://github.com/apache/spark/blob/36e3b6d5ad17ecc6afee7dae10b45b985cc72666/python/pyspark/sql/connect/client/core.py#L1547-L1548))
and `config`
([L2051-L2052](https://github.com/apache/spark/blob/36e3b6d5ad17ecc6afee7dae10b45b985cc72666/python/pyspark/sql/connect/client/core.py#L2051-L2052)).
The `""` comes from `_verify_response_integrity`, which stores whatever the
first response carries
([L2555-L2557](https://github.com/apache/spark/blob/36e3b6d5ad17ecc6afee7dae10b45b985cc72666/python/pyspark/sql/connect/client/core.py#L2555-L2557)).
The stored `""` also causes a problem outside the exit hook. It was there
before this PR, so it may be better as a separate JIRA. After a no-op
`ReleaseSession`, the next `Config` or `ExecutePlan` sends
`client_observed_server_side_session_id=""`. The server creates a new session,
and then `validateSessionId` throws `INVALID_HANDLE.SESSION_CHANGED`
([L93-L105](https://github.com/apache/spark/blob/36e3b6d5ad17ecc6afee7dae10b45b985cc72666/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala#L93-L105)).
The client marks itself closed, and the new server session stays until the
idle timeout.
If `_verify_response_integrity` only stored a non-empty ID,
`_server_session_id` would be either `None` or a real ID. Then this check would
not need the special case, and the `""` subtests would not be needed either:
```python
elif response.server_side_session_id:
# Update the server side session ID.
self._server_session_id = response.server_side_session_id
```
##########
python/pyspark/sql/connect/client/core.py:
##########
@@ -2202,6 +2203,7 @@ def release_session(self) -> None:
timeout=self._rpc_deadlines.release_session,
)
self._verify_response_integrity(resp)
+ self._session_released = True
Review Comment:
**1. [Low] A no-op `ReleaseSession` also marks the session as released**
`release_session()` sets `_session_released` after any successful response,
including the case where the server had no session for this ID. In that case
the server returns an empty `server_side_session_id`
([L37-L38](https://github.com/apache/spark/blob/36e3b6d5ad17ecc6afee7dae10b45b985cc72666/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectReleaseSessionHandler.scala#L37-L38)).
`closeSession` also does not add the key to `closedSessionsCache`, because
`removeSessionHolder` only does that for an existing holder
([L237-L248](https://github.com/apache/spark/blob/36e3b6d5ad17ecc6afee7dae10b45b985cc72666/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala#L237-L248)).
So the session ID can still be used, but `_on_exit()` now returns early for
good.
For example, with `SPARK_CONNECT_RELEASE_SESSION_ON_EXIT=true`:
```python
spark.client.release_session() # no server session yet, so a no-op on the
server
spark.addArtifacts("dep.py", pyfile=True) # the server creates a new session
# at exit, _on_exit() returns early and sends no ReleaseSession
```
`ArtifactManager` sends no observed ID, so the server creates the session
silently
([L56-L63](https://github.com/apache/spark/blob/36e3b6d5ad17ecc6afee7dae10b45b985cc72666/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectAddArtifactsHandler.scala#L56-L63)).
The session and the artifact then stay on the server until the idle timeout.
This is the artifact-upload case that the description says should still be
released at exit.
Could the flag be set only when the server actually closed a session, that
is, when the response has a non-empty `server_side_session_id`?
```python
self._verify_response_integrity(resp)
if resp.server_side_session_id:
self._session_released = True
```
With this, the "released before use" subtest would expect one
`ReleaseSession` when release-on-exit is enabled. That RPC is a no-op on the
server.
##########
python/pyspark/sql/connect/client/core.py:
##########
@@ -2604,14 +2606,13 @@ def _delete_ml_cache(self, cache_ids: List[str],
evict_only: bool = False) -> Li
return []
def _on_exit(self) -> None:
- # If the client has already been explicitly closed, skip all cleanup
RPCs.
- # The server-side resources were released by close(); reissuing them
here
- # is wasted work and, if the server has since become unreachable, can
- # block process exit on the gRPC call.
- if self._closed:
+ # Skip cleanup RPCs if the session is already closed or released.
Review Comment:
**4. [Nit] The new comment drops the reason, and the early return also skips
`close()`**
The removed comment explained why a closed client must not send exit RPCs:
the work is wasted, and with an unreachable server it can block process exit
(#56140). The new comment only says what the code does. Could it keep the
reason?
Also, the early return skips more than RPCs. With
`SPARK_CONNECT_RELEASE_SESSION_ON_EXIT=true`, a client that called
`release_session()` but not `close()` used to run `close()` at exit. `close()`
waits for `_release_futures`, shuts down the reattach thread pool, and closes
the channel
([L1490-L1497](https://github.com/apache/spark/blob/36e3b6d5ad17ecc6afee7dae10b45b985cc72666/python/pyspark/sql/connect/client/core.py#L1490-L1497)).
Now it does not run. That is probably fine at process exit, but the comment
could say "skip all exit cleanup" instead of "skip cleanup RPCs".
##########
python/pyspark/sql/connect/client/core.py:
##########
@@ -2604,14 +2606,13 @@ def _delete_ml_cache(self, cache_ids: List[str],
evict_only: bool = False) -> Li
return []
def _on_exit(self) -> None:
- # If the client has already been explicitly closed, skip all cleanup
RPCs.
- # The server-side resources were released by close(); reissuing them
here
- # is wasted work and, if the server has since become unreachable, can
- # block process exit on the gRPC call.
- if self._closed:
+ # Skip cleanup RPCs if the session is already closed or released.
+ if self._closed or self._session_released:
Review Comment:
**3. [Nit] Unregistering the exit hook could replace the new flag**
The PR adds a state that `_on_exit()` checks at exit. Another way is to
remove the hook when the client reaches that state: call
`atexit.unregister(self._on_exit)` in `close()` and after a successful
`release_session()`. `atexit.unregister` compares with `==`, so a new bound
method of the same client matches the registered one.
This would also fix an older issue. The atexit registry keeps a strong
reference to every `SparkConnectClient`, with its channel and
`ArtifactManager`, until the process exits, even after `close()`. An app that
calls `newSession()` or `clone()` many times keeps all of these clients. The
`_closed` check would stay, because `_closed` is also set on `SESSION_CHANGED`.
This is only an alternative. The flag works as it is.
##########
python/pyspark/sql/tests/connect/client/test_client.py:
##########
@@ -254,10 +254,69 @@ def GetStatus(self, req: proto.GetStatusRequest,
metadata, timeout=None):
resp.operation_statuses.append(op_status)
return resp
+ class ExitRecordingStub:
+ """Records every RPC by name, so a test can assert which RPCs the
client sent.
+
+ Answers the RPCs the exit hook may send; any other RPC is recorded and
then fails.
+ """
+
+ def __init__(self, session_id: str, release_server_side_session_id:
str = ""):
+ self._session_id = session_id
+ # ReleaseSession carries no server-side ID when the session never
existed.
+ self._release_server_side_session_id =
release_server_side_session_id
+ self.calls = []
+
+ def Config(self, req, metadata, timeout=None):
+ self.calls.append("Config")
+ resp = proto.ConfigResponse(
+ session_id=self._session_id,
server_side_session_id="server-session-id"
+ )
+ for key in req.operation.get.keys:
+ resp.pairs.add(key=key)
+ return resp
+
+ def ExecutePlan(self, req, metadata, timeout=None):
+ self.calls.append("ExecutePlan")
+ return [
+ proto.ExecutePlanResponse(
+ session_id=self._session_id,
+ server_side_session_id="server-session-id",
+ operation_id=req.operation_id,
+ )
+ ]
+
+ def ReleaseSession(self, req, metadata, timeout=None):
+ self.calls.append("ReleaseSession")
+ return proto.ReleaseSessionResponse(
+ session_id=self._session_id,
+ server_side_session_id=self._release_server_side_session_id,
+ )
+
+ def __getattr__(self, name):
+ if name.startswith("__"):
+ raise AttributeError(name)
+
+ def unexpected_rpc(*args, **kwargs):
+ self.calls.append(name)
+ raise RuntimeError(f"Unexpected RPC: {name}")
+
+ return unexpected_rpc
+
# The _cleanup_ml_cache invocation will hang in this test (no valid spark
cluster)
# and it blocks the test process exiting because it is registered as the
atexit handler
# in `SparkConnectClient` constructor. To bypass the issue, patch the
method in the test.
- SparkConnectClient._cleanup_ml_cache = lambda _: None
+ # The real method stays reachable through __wrapped__ because the test
runner can import
+ # this module twice (as __main__ and by name), and the second import must
not capture the
+ # no-op.
+ _real_cleanup_ml_cache = getattr(
+ SparkConnectClient._cleanup_ml_cache, "__wrapped__",
SparkConnectClient._cleanup_ml_cache
+ )
+
+ def _skip_cleanup_ml_cache(self):
Review Comment:
**9. [Nit] `functools.wraps` and `inspect.unwrap` already do this**
The standard library has this convention. `functools.wraps(real)` sets
`__wrapped__` and also keeps `__name__` and `__qualname__`, and
`inspect.unwrap()` follows `__wrapped__`. With `functools` and `inspect`
imported:
```python
_real_cleanup_ml_cache =
inspect.unwrap(SparkConnectClient._cleanup_ml_cache)
@functools.wraps(_real_cleanup_ml_cache)
def _skip_cleanup_ml_cache(self):
pass
SparkConnectClient._cleanup_ml_cache = _skip_cleanup_ml_cache
```
##########
python/pyspark/sql/tests/connect/client/test_client.py:
##########
@@ -964,6 +1143,7 @@ def GetStatus(self, req, metadata, timeout=None):
client.release_session()
self.assertEqual(mock.captured_timeouts["ReleaseSession"], 44.0)
+ self.assertTrue(client._session_released)
Review Comment:
**7. [Nit] No test covers a failed `release_session()`**
This checks that a successful `release_session()` sets `_session_released`.
No test checks the other side: when `ReleaseSession` fails, the flag stays
`False` and `_on_exit()` still cleans up. If the assignment moved before
`_verify_response_integrity(resp)`, or into a `finally` block, every test would
still pass.
Could a test make `ReleaseSession` raise a non-retryable error, and assert
that `_session_released` stays `False` and that `_on_exit()` still sends its
RPCs?
##########
python/pyspark/sql/tests/connect/client/test_client.py:
##########
@@ -254,10 +254,69 @@ def GetStatus(self, req: proto.GetStatusRequest,
metadata, timeout=None):
resp.operation_statuses.append(op_status)
return resp
+ class ExitRecordingStub:
+ """Records every RPC by name, so a test can assert which RPCs the
client sent.
+
+ Answers the RPCs the exit hook may send; any other RPC is recorded and
then fails.
+ """
+
+ def __init__(self, session_id: str, release_server_side_session_id:
str = ""):
+ self._session_id = session_id
+ # ReleaseSession carries no server-side ID when the session never
existed.
+ self._release_server_side_session_id =
release_server_side_session_id
+ self.calls = []
+
+ def Config(self, req, metadata, timeout=None):
+ self.calls.append("Config")
+ resp = proto.ConfigResponse(
+ session_id=self._session_id,
server_side_session_id="server-session-id"
+ )
+ for key in req.operation.get.keys:
+ resp.pairs.add(key=key)
+ return resp
+
+ def ExecutePlan(self, req, metadata, timeout=None):
+ self.calls.append("ExecutePlan")
+ return [
+ proto.ExecutePlanResponse(
+ session_id=self._session_id,
+ server_side_session_id="server-session-id",
+ operation_id=req.operation_id,
+ )
+ ]
+
+ def ReleaseSession(self, req, metadata, timeout=None):
+ self.calls.append("ReleaseSession")
+ return proto.ReleaseSessionResponse(
+ session_id=self._session_id,
+ server_side_session_id=self._release_server_side_session_id,
+ )
+
+ def __getattr__(self, name):
+ if name.startswith("__"):
+ raise AttributeError(name)
+
+ def unexpected_rpc(*args, **kwargs):
+ self.calls.append(name)
+ raise RuntimeError(f"Unexpected RPC: {name}")
+
+ return unexpected_rpc
+
# The _cleanup_ml_cache invocation will hang in this test (no valid spark
cluster)
# and it blocks the test process exiting because it is registered as the
atexit handler
# in `SparkConnectClient` constructor. To bypass the issue, patch the
method in the test.
- SparkConnectClient._cleanup_ml_cache = lambda _: None
+ # The real method stays reachable through __wrapped__ because the test
runner can import
+ # this module twice (as __main__ and by name), and the second import must
not capture the
+ # no-op.
+ _real_cleanup_ml_cache = getattr(
Review Comment:
**6. [Low] `_real_cleanup_ml_cache` breaks if another module patches the
method first**
This finds the real method only if this module is the first to patch the
class. `test_client_call_stack_trace.py` patches it with a plain lambda that
has no `__wrapped__`
([L33](https://github.com/apache/spark/blob/36e3b6d5ad17ecc6afee7dae10b45b985cc72666/python/pyspark/sql/tests/connect/client/test_client_call_stack_trace.py#L33)).
If that module is imported first in the same process, `_real_cleanup_ml_cache`
becomes that lambda. For example:
```bash
python -m pytest
python/pyspark/sql/tests/connect/client/test_client_call_stack_trace.py \
python/pyspark/sql/tests/connect/client/test_client.py
```
Then the control check in `test_on_exit_sends_no_rpc_without_live_session`
gets `[]` and fails, for a reason that has nothing to do with `_on_exit()`.
`python/run-tests` runs each module in its own process, so CI is not
affected. Could both modules use one shared helper, for example in
`pyspark/testing/connectutils.py`, that patches the method once and keeps the
original?
##########
python/pyspark/sql/tests/connect/client/test_client.py:
##########
@@ -745,6 +878,52 @@ def mock_close():
self.assertEqual(call_tracker["release_session"], 1)
self.assertEqual(call_tracker["close"], 1)
+ def test_on_exit_sends_no_rpc_without_live_session(self):
+ """Test _on_exit sends no RPC for a session that was never established
or was released.
+
+ RPCs are recorded at the stub, so any new exit-time RPC is covered,
not only the ones
+ _on_exit sends today.
+ """
+
+ def exit_rpcs(observed_id, release_on_exit, release_first=False,
release_id=""):
+ client = SparkConnectClient("sc://foo/",
use_reattachable_execute=False)
+ # Use the real ML cache cleanup, which this module patches out,
and keep the atexit
+ # hook from running it again at interpreter exit.
+ client._cleanup_ml_cache = _real_cleanup_ml_cache.__get__(client)
+ self.addCleanup(setattr, client, "_closed", True)
Review Comment:
**10. [Nit] The cleanup sets `_closed` but does not close the channel**
Setting `_closed` keeps the atexit hook quiet, but the gRPC channel of each
client stays open until the process exits. This test creates nine clients.
Other tests in this file call `client.close()` for the same reason
([L570-L571](https://github.com/apache/spark/blob/36e3b6d5ad17ecc6afee7dae10b45b985cc72666/python/pyspark/sql/tests/connect/client/test_client.py#L570-L571)).
Could this be `self.addCleanup(client.close)`? It also sets `_closed`, and
as far as I can tell it is safe to call after `_on_exit()` has already closed
the client.
##########
python/pyspark/sql/tests/connect/client/test_client.py:
##########
@@ -623,11 +683,80 @@ def mock_close():
self.assertEqual(call_tracker["release_session"], 1)
self.assertEqual(call_tracker["close"], 1)
+ def test_on_exit_does_not_call_when_session_released(self):
Review Comment:
**8. [Nit] `test_on_exit_does_not_call_when_session_released` repeats the
closed-client test**
This test is the same as `test_on_exit_does_not_call_when_already_closed`
([L777-L803](https://github.com/apache/spark/blob/36e3b6d5ad17ecc6afee7dae10b45b985cc72666/python/pyspark/sql/tests/connect/client/test_client.py#L777-L803))
except for the flag it sets. The "released after use" subtest of
`test_on_exit_sends_no_rpc_without_live_session` also checks the same thing
through the real `release_session()`.
Could the two mock-based tests become one test with a `subTest` for each
flag?
##########
python/pyspark/sql/connect/client/core.py:
##########
@@ -2604,14 +2606,13 @@ def _delete_ml_cache(self, cache_ids: List[str],
evict_only: bool = False) -> Li
return []
def _on_exit(self) -> None:
- # If the client has already been explicitly closed, skip all cleanup
RPCs.
- # The server-side resources were released by close(); reissuing them
here
- # is wasted work and, if the server has since become unreachable, can
- # block process exit on the gRPC call.
- if self._closed:
+ # Skip cleanup RPCs if the session is already closed or released.
+ if self._closed or self._session_released:
return
- self._cleanup_ml_cache()
+ # No server-side session ID (None or "") means no ML cache to clean.
+ if self._server_session_id:
+ self._cleanup_ml_cache()
Review Comment:
**5. [Nit] Sessions that never used ML still send `clean_cache` at exit**
This check skips the cleanup only for a client that never saw a server-side
session ID. A client that ran any query has seen one, so a SQL-only app still
sends `ml_command.clean_cache` at exit. With a slow or unreachable server, exit
can still wait for that RPC's deadline and retries. The description gives this
delay as a reason for the PR, but it remains for the common case.
A narrower condition would be "this client sent an ML command or an ML
relation". I am not sure it is worth it, because ML relations such as a model
summary fetch can also add entries to the cache, so such a flag must cover them
too. This is out of scope here. Could the description say that the exit delay
remains for sessions that ran queries?
##########
python/pyspark/sql/tests/connect/client/test_client.py:
##########
@@ -745,6 +878,52 @@ def mock_close():
self.assertEqual(call_tracker["release_session"], 1)
self.assertEqual(call_tracker["close"], 1)
+ def test_on_exit_sends_no_rpc_without_live_session(self):
Review Comment:
**11. [Nit] The test list and the test count in the description are out of
date**
The description says "All 44 tests passed" and lists two new tests. This
module now has 45 tests. The description does not mention this test or the
`__wrapped__` change in the module-level patch, and both came in the last two
commits. Could the description list this test as well? It is the one that would
catch a new exit-time RPC.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
