diff --git a/python/pyspark/sql/connect/local_server.py b/python/pyspark/sql/connect/local_server.py index 4f06d884aff49..4aa88a400487e 100644 --- a/python/pyspark/sql/connect/local_server.py +++ b/python/pyspark/sql/connect/local_server.py @@ -75,6 +75,22 @@ def _pid_alive(pid: int) -> bool: return True +def runtime_dir() -> str: + """Return the private per-user directory holding local-server state.""" + path = os.path.join(tempfile.gettempdir(), "spark-connect-{}".format(getpass.getuser())) + try: + # exist_ok also covers two first runs racing to create the directory; chmod + # re-asserts 0700 and fails if another user owns the path. + os.makedirs(path, mode=0o700, exist_ok=True) + os.chmod(path, 0o700) + except OSError as e: + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_RUNTIME_DIR_UNAVAILABLE", + messageParameters={"path": path}, + ) from e + return path + + class Discovery: """Reads and writes the discovery file recording the persistent local server. @@ -82,26 +98,11 @@ class Discovery: ``SPARK_LOCAL_CONNECT_DISCOVERY`` points; the daemon's pid file and logs sit next to it. """ - @staticmethod - def _runtime_dir() -> str: - path = os.path.join(tempfile.gettempdir(), "spark-connect-{}".format(getpass.getuser())) - try: - # exist_ok also covers two first runs racing to create the directory; chmod - # re-asserts 0700 and fails if another user owns the path. - os.makedirs(path, mode=0o700, exist_ok=True) - os.chmod(path, 0o700) - except OSError as e: - raise PySparkRuntimeError( - errorClass="LOCAL_CONNECT_RUNTIME_DIR_UNAVAILABLE", - messageParameters={"path": path}, - ) from e - return path - def __init__(self, path: Optional[str] = None): self.path = os.path.abspath( path or os.environ.get("SPARK_LOCAL_CONNECT_DISCOVERY") - or os.path.join(self._runtime_dir(), "connect-local.json") + or os.path.join(runtime_dir(), "connect-local.json") ) self._lock_file: Optional[TextIO] = None @@ -219,12 +220,33 @@ def is_reusable(self) -> bool: def reuse_or_start(self, master: str, opts: Dict[str, Any]) -> str: if not self.is_reusable(): - ServerLauncher(master, opts, self._discovery).launch() - self._reload() + self.start(master, opts) assert self.token is not None os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = self.token return self.url + def start( + self, + master: str, + opts: Dict[str, Any], + *, + use_ephemeral_port: bool = False, + seed_conf: Optional[Dict[str, Any]] = None, + ) -> None: + """Start this server and reload its discovery record. + + Callers starting isolated daemons can request an ephemeral port and provide a + precomputed startup configuration while sharing the standard launch path. + """ + ServerLauncher( + master, + opts, + self._discovery, + use_ephemeral_port=use_ephemeral_port, + seed_conf=seed_conf, + ).launch() + self._reload() + def stop(self) -> bool: stopped = False if self.pid is not None: @@ -237,17 +259,58 @@ def stop(self) -> bool: return stopped +def _strip_launcher_conf(conf: Dict[str, Any]) -> Dict[str, Any]: + """Drop keys the launcher sets itself (master, binding port, auth token) and the + ``spark.local.connect.*`` opt-in keys, returning a new dict. Idempotent, so it is safe to + apply to a conf that is already sanitized. + """ + stripped = dict(conf) + for k in list(stripped): + if k in ( + "spark.remote", + "spark.api.mode", + "spark.master", + "spark.connect.authenticate.token", + "spark.connect.grpc.binding.port", + ) or k.startswith("spark.local.connect."): + stripped.pop(k) + return stripped + + +def startup_seed_conf(opts: Dict[str, Any]) -> Dict[str, Any]: + """Compute startup confs using the same merge as the in-process server path, then strip + the keys the launcher manages (see ``_strip_launcher_conf``). + """ + conf: Dict[str, Any] = {} + for i in range(int(os.environ.get("PYSPARK_REMOTE_INIT_CONF_LEN", "0"))): + conf = json.loads(os.environ["PYSPARK_REMOTE_INIT_CONF_{}".format(i)]) + conf.update(opts) + return _strip_launcher_conf(conf) + + class ServerLauncher: """Starts a persistent local server via ``sbin/start-connect-server.sh`` and waits until it accepts connections. Callers must hold a ``Discovery`` context. + + ``use_ephemeral_port`` and ``seed_conf`` allow other managed local servers to share this + launch path without duplicating its process and readiness handling. """ _READY_TIMEOUT = 120 - def __init__(self, master: str, opts: Dict[str, Any], discovery: Discovery): + def __init__( + self, + master: str, + opts: Dict[str, Any], + discovery: Discovery, + use_ephemeral_port: bool = False, + seed_conf: Optional[Dict[str, Any]] = None, + ): self._master = master self._opts = opts self._discovery = discovery + self._use_ephemeral_port = use_ephemeral_port + self._seed_override = seed_conf self._log_dir = os.path.join(discovery.directory, "logs") def launch(self) -> None: @@ -269,12 +332,12 @@ def _token(self) -> str: ) def _pick_port(self) -> int: - """Under SPARK_TESTING always use an OS-assigned free port so suites can run in - parallel; otherwise honor the configured/default port, falling back to a free one if - another process holds it. (A live stale server of ours also holds the port, but that - start fails later at spark-daemon.sh's pid-file check regardless of port.) The sbin - script cannot report an ephemeral port back, so the free port is picked and released - here, with a small race until the server binds it. + """Use an OS-assigned free port when requested or under SPARK_TESTING so suites can + run in parallel. Otherwise honor the configured/default port, falling back to a free + one if another process holds it. (A live stale server of ours also holds the port, but + that start fails later at spark-daemon.sh's pid-file check regardless of port.) The + sbin script cannot report an ephemeral port back, so the free port is picked and + released here, with a small race until the server binds it. """ def free_port() -> int: @@ -282,7 +345,7 @@ def free_port() -> int: sock.bind(("localhost", 0)) return sock.getsockname()[1] - if "SPARK_TESTING" in os.environ: + if self._use_ephemeral_port or "SPARK_TESTING" in os.environ: return free_port() from pyspark.sql.connect.client import DefaultChannelBuilder @@ -297,26 +360,19 @@ def free_port() -> int: return free_port() def _seed_conf(self) -> Dict[str, Any]: - """Startup confs for the new server, merged like the in-process - ``_start_connect_server`` merges ``PYSPARK_REMOTE_INIT_CONF_*`` with the builder - opts, minus the keys the launcher sets itself and the ``spark.local.connect.*`` - opt-in keys. Only the run that starts the server can seed static confs; later runs - find the JVM already warm. + """Startup confs for the new server, minus the keys the launcher sets itself and the + ``spark.local.connect.*`` opt-in keys. Only the run that starts the server can seed + static confs; later runs find the JVM already warm. + + With no ``seed_conf`` override, this merges ``PYSPARK_REMOTE_INIT_CONF_*`` with the + builder opts like the in-process ``_start_connect_server`` does. When an override is + given, it is used verbatim instead of the merge. Either way the result is run through + ``_strip_launcher_conf``, so the launcher-managed keys never reach + ``--properties-file`` even if a caller passes raw opts as the override. """ - conf: Dict[str, Any] = {} - for i in range(int(os.environ.get("PYSPARK_REMOTE_INIT_CONF_LEN", "0"))): - conf = json.loads(os.environ["PYSPARK_REMOTE_INIT_CONF_{}".format(i)]) - conf.update(self._opts) - for k in list(conf): - if k in ( - "spark.remote", - "spark.api.mode", - "spark.master", - "spark.connect.authenticate.token", - "spark.connect.grpc.binding.port", - ) or k.startswith("spark.local.connect."): - conf.pop(k) - return conf + if self._seed_override is not None: + return _strip_launcher_conf(self._seed_override) + return startup_seed_conf(self._opts) @contextlib.contextmanager def _seed_properties_file(self) -> Iterator[Optional[str]]: diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server.py b/python/pyspark/sql/tests/connect/test_connect_local_server.py index eba26972ae8c3..ae09eb04646f5 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server.py @@ -104,6 +104,26 @@ def _discovered_server(self) -> "LocalConnectServer": with Discovery() as discovery: return LocalConnectServer(discovery) + def _launcher_discovery(self): + # A stand-in Discovery for ServerLauncher unit tests: only its directory is read + # (for the log dir and the seed properties file), so point it at the temp dir. + from unittest import mock + + discovery = mock.Mock() + discovery.directory = self._tmpdir + return discovery + + @contextlib.contextmanager + def _without_spark_testing(self): + # _pick_port's ephemeral branch is a no-op when SPARK_TESTING is set (as it is under + # the test runner), so drop it to exercise the production behavior. + saved = os.environ.pop("SPARK_TESTING", None) + try: + yield + finally: + if saved is not None: + os.environ["SPARK_TESTING"] = saved + def test_discovery_location(self) -> None: self.assertEqual(Discovery().path, self._discovery_path) # Without the override the file lives in a per-user 0700 dir under the temp dir. @@ -114,6 +134,170 @@ def test_discovery_location(self) -> None: self.assertIn("spark-connect-{}".format(getpass.getuser()), default.directory) self.assertEqual(os.stat(default.directory).st_mode & 0o777, 0o700) + def test_startup_seed_conf(self) -> None: + from unittest import mock + + initial = { + "spark.sql.shuffle.partitions": "8", + "spark.master": "local[1]", + } + opts = { + "spark.sql.warehouse.dir": os.path.join(self._tmpdir, "warehouse"), + "spark.local.connect.reuse": "true", + "spark.connect.grpc.binding.port": "0", + } + env = { + "PYSPARK_REMOTE_INIT_CONF_LEN": "1", + "PYSPARK_REMOTE_INIT_CONF_0": json.dumps(initial), + } + with mock.patch.dict(os.environ, env): + self.assertEqual( + local_server.startup_seed_conf(opts), + { + "spark.sql.shuffle.partitions": "8", + "spark.sql.warehouse.dir": opts["spark.sql.warehouse.dir"], + }, + ) + + def test_start_delegates_launch_options(self) -> None: + from unittest import mock + + discovery = mock.Mock() + discovery.load.side_effect = [ + None, + { + "host": "localhost", + "port": 15002, + "token": "t", + "pid": os.getpid(), + "spark_version": __version__, + }, + ] + server = LocalConnectServer(discovery) + seed_conf = {"spark.sql.shuffle.partitions": "4"} + with mock.patch.object(local_server, "ServerLauncher") as launcher: + server.start( + "local[2]", + {"spark.local.connect.reuse": "true"}, + use_ephemeral_port=True, + seed_conf=seed_conf, + ) + + launcher.assert_called_once_with( + "local[2]", + {"spark.local.connect.reuse": "true"}, + discovery, + use_ephemeral_port=True, + seed_conf=seed_conf, + ) + launcher.return_value.launch.assert_called_once_with() + self.assertEqual(server.port, 15002) + + def test_pick_port_uses_ephemeral_port_when_requested(self) -> None: + # This is the production path for pool attendants, which run without SPARK_TESTING. + # A non-integer configured port would raise int() in the configured/default branch; + # the ephemeral branch never reads it, so returning a clean OS-assigned port proves + # the free-port path was taken even with SPARK_TESTING unset. + launcher = local_server.ServerLauncher( + "local[2]", + {"spark.local.connect.server.port": "not-a-port"}, + self._launcher_discovery(), + use_ephemeral_port=True, + ) + with self._without_spark_testing(): + port = launcher._pick_port() + self.assertGreater(port, 0) + + def test_pick_port_honors_configured_port_without_testing(self) -> None: + # With neither the ephemeral flag nor SPARK_TESTING, a free configured port is used + # as-is rather than replaced by an OS-assigned one. + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("localhost", 0)) + free = sock.getsockname()[1] + launcher = local_server.ServerLauncher( + "local[2]", + {"spark.local.connect.server.port": str(free)}, + self._launcher_discovery(), + use_ephemeral_port=False, + ) + with self._without_spark_testing(): + self.assertEqual(launcher._pick_port(), free) + + def test_seed_conf_override_is_used_and_sanitized(self) -> None: + # The override is taken verbatim except for launcher-managed keys, which are stripped + # so raw builder opts passed as a seed cannot land in --properties-file. + launcher = local_server.ServerLauncher( + "local[2]", + {}, + self._launcher_discovery(), + seed_conf={ + "spark.sql.shuffle.partitions": "4", + "spark.master": "local[9]", + "spark.local.connect.reuse": "true", + }, + ) + self.assertEqual(launcher._seed_conf(), {"spark.sql.shuffle.partitions": "4"}) + + def test_seed_conf_empty_override_does_not_fall_through_to_env(self) -> None: + from unittest import mock + + # Load-bearing for the pool attendant: an empty override means "seed nothing", and + # must not silently pick up PYSPARK_REMOTE_INIT_CONF_* the way opts=None would. + env = { + "PYSPARK_REMOTE_INIT_CONF_LEN": "1", + "PYSPARK_REMOTE_INIT_CONF_0": json.dumps({"spark.sql.shuffle.partitions": "8"}), + } + launcher = local_server.ServerLauncher( + "local[2]", {}, self._launcher_discovery(), seed_conf={} + ) + with mock.patch.dict(os.environ, env): + self.assertEqual(launcher._seed_conf(), {}) + + def test_seed_conf_none_override_merges_env_and_opts(self) -> None: + from unittest import mock + + # No override: the env-plus-opts merge (minus launcher-managed keys) is used. + env = { + "PYSPARK_REMOTE_INIT_CONF_LEN": "1", + "PYSPARK_REMOTE_INIT_CONF_0": json.dumps({"spark.sql.shuffle.partitions": "8"}), + } + launcher = local_server.ServerLauncher( + "local[2]", + {"spark.sql.warehouse.dir": os.path.join(self._tmpdir, "wh")}, + self._launcher_discovery(), + seed_conf=None, + ) + with mock.patch.dict(os.environ, env): + self.assertEqual( + launcher._seed_conf(), + { + "spark.sql.shuffle.partitions": "8", + "spark.sql.warehouse.dir": os.path.join(self._tmpdir, "wh"), + }, + ) + + def test_seed_properties_file_reflects_seed_conf(self) -> None: + # An empty seed yields no properties file, so start-connect-server.sh gets no + # --properties-file; a non-empty seed writes a 0600 file with the seeded confs. + launcher = local_server.ServerLauncher( + "local[2]", {}, self._launcher_discovery(), seed_conf={} + ) + with launcher._seed_properties_file() as path: + self.assertIsNone(path) + + launcher = local_server.ServerLauncher( + "local[2]", + {}, + self._launcher_discovery(), + seed_conf={"spark.sql.shuffle.partitions": "4"}, + ) + with launcher._seed_properties_file() as path: + self.assertIsNotNone(path) + self.assertEqual(os.stat(path).st_mode & 0o777, 0o600) + with open(path) as f: + contents = f.read() + self.assertIn("spark.sql.shuffle.partitions=4", contents) + def test_discovery_roundtrip(self) -> None: with Discovery() as discovery: saved = self._server(port=15002)