ericm-db commented on code in PR #56907: URL: https://github.com/apache/spark/pull/56907#discussion_r3540191039
########## python/pyspark/sql/tests/connect/test_connect_local_server.py: ########## @@ -0,0 +1,353 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import json +import os +import shutil +import signal +import socket +import sys +import tempfile +import time +import unittest + +from pyspark.util import is_remote_only +from pyspark.testing.connectutils import should_test_connect, connect_requirement_message + +if should_test_connect: + from pyspark.sql.connect.session import SparkSession as RemoteSparkSession + from pyspark.version import __version__ + + [email protected]( + not should_test_connect or is_remote_only(), + connect_requirement_message or "Requires JVM access to start a local Connect server", +) +class LocalConnectServerReuseTests(unittest.TestCase): + """Tests for the opt-in persistent local Spark Connect server (SPARK_LOCAL_CONNECT_REUSE).""" + + def setUp(self) -> None: + # Point discovery at a throwaway path and remember the env we override, so each test starts + # from a clean slate and the real ~/.spark/connect-local.json is never touched. + self._tmpdir = tempfile.mkdtemp() + self._discovery = os.path.join(self._tmpdir, "connect-local.json") + self._saved_env = { + k: os.environ.get(k) + for k in ("SPARK_LOCAL_CONNECT_DISCOVERY", "SPARK_CONNECT_AUTHENTICATE_TOKEN") + } + os.environ["SPARK_LOCAL_CONNECT_DISCOVERY"] = self._discovery + + def tearDown(self) -> None: + try: + # Only stop a real, separately-spawned server. The discovery-logic unit tests fabricate + # discovery files that point at this very process, which must never be signalled. + disc = RemoteSparkSession._read_local_connect_discovery() + if disc is not None and disc.get("pid") != os.getpid(): + port = int(disc["port"]) + RemoteSparkSession._stop_local_connect_server() + # _stop_local_connect_server only signals the daemon and returns; wait for the JVM + # to actually release the port so the next test starts from a clean slate. + self._wait_port_closed(disc["host"], port) + finally: + for k, v in self._saved_env.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + # Remove the whole scratch dir: besides the discovery file it may hold a .lock file, + # seed-conf temp files, and a seeded warehouse directory. + shutil.rmtree(self._tmpdir, ignore_errors=True) + + # -- discovery / reuse-decision logic (no real server) ---------------------------------------- + + def test_discovery_path_honors_override(self) -> None: + self.assertEqual(RemoteSparkSession._local_connect_discovery_path(), self._discovery) + os.environ.pop("SPARK_LOCAL_CONNECT_DISCOVERY") + self.assertTrue( + RemoteSparkSession._local_connect_discovery_path().endswith( + os.path.join(".spark", "connect-local.json") + ) + ) + + def test_read_discovery_missing_or_malformed(self) -> None: + self.assertIsNone(RemoteSparkSession._read_local_connect_discovery()) + with open(self._discovery, "w") as f: + f.write("not json") + self.assertIsNone(RemoteSparkSession._read_local_connect_discovery()) + with open(self._discovery, "w") as f: + json.dump({"host": "localhost"}, f) # missing required keys + self.assertIsNone(RemoteSparkSession._read_local_connect_discovery()) + + def _write_discovery(self, **overrides) -> dict: + disc = { + "host": "localhost", + "port": 0, + "token": "t", + "pid": os.getpid(), + "spark_version": __version__, + } + disc.update(overrides) + with open(self._discovery, "w") as f: + json.dump(disc, f) + return disc + + def test_not_reusable_on_version_mismatch(self) -> None: + disc = self._write_discovery(spark_version="0.0.0-not-this-build") + self.assertFalse(RemoteSparkSession._local_connect_server_is_reusable(disc)) + + def test_not_reusable_on_dead_pid(self) -> None: + # PID 2**31 - 1 is effectively guaranteed not to exist. + disc = self._write_discovery(pid=2**31 - 1, port=1) + self.assertFalse(RemoteSparkSession._local_connect_server_is_reusable(disc)) + + def test_reusable_when_alive_and_listening(self) -> None: + # A live listening socket owned by this (alive) process with a matching version is reusable. + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + listener.bind(("localhost", 0)) + listener.listen(1) + port = listener.getsockname()[1] + disc = self._write_discovery(port=port) + self.assertTrue(RemoteSparkSession._local_connect_server_is_reusable(disc)) + finally: + listener.close() + # Once the socket is closed the port is no longer reachable, so it is not reusable. + self.assertFalse(RemoteSparkSession._local_connect_server_is_reusable(disc)) + + def test_stop_when_no_server_is_safe(self) -> None: + self.assertFalse(RemoteSparkSession._stop_local_connect_server()) + + def test_reuse_from_discovery_none_when_absent(self) -> None: + self.assertIsNone(RemoteSparkSession._reuse_from_discovery()) + + def test_local_port_available(self) -> None: + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + listener.bind(("localhost", 0)) + listener.listen(1) + taken = listener.getsockname()[1] + self.assertFalse(RemoteSparkSession._local_port_available(taken)) + finally: + listener.close() + # The port is free again once the listener is closed. + self.assertTrue(RemoteSparkSession._local_port_available(taken)) + + def test_server_conf_seeds_user_confs_and_drops_control_keys(self) -> None: + """_local_connect_server_conf keeps user startup confs but not keys the daemon controls.""" + opts = { + "spark.sql.warehouse.dir": "/tmp/wh", + "spark.jars.packages": "org.example:lib:1.0", + "spark.remote": "local[*]", + "spark.master": "local[*]", + "spark.connect.authenticate.token": "secret", + "spark.connect.grpc.binding.port": "15002", + "spark.local.connect.reuse": "true", + "spark.local.connect.server.port": "15002", + "spark.local.connect.server.idleTimeout": "60", + } + conf = RemoteSparkSession._local_connect_server_conf(opts) + self.assertEqual(conf.get("spark.sql.warehouse.dir"), "/tmp/wh") + self.assertEqual(conf.get("spark.jars.packages"), "org.example:lib:1.0") + for dropped in ( + "spark.remote", + "spark.master", + "spark.connect.authenticate.token", + "spark.connect.grpc.binding.port", + "spark.local.connect.reuse", + "spark.local.connect.server.port", + "spark.local.connect.server.idleTimeout", + ): + self.assertNotIn(dropped, conf) + + def test_start_lock_roundtrip(self) -> None: + """Acquiring and releasing the start-up lock creates the lock file and does not error.""" + fd = RemoteSparkSession._acquire_local_connect_start_lock() Review Comment: Yup added this to the docs -- 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]
