ericm-db commented on code in PR #56907:
URL: https://github.com/apache/spark/pull/56907#discussion_r3547854202
##########
python/pyspark/sql/connect/session.py:
##########
@@ -1247,6 +1249,372 @@ def _start_connect_server(master: str, opts: Dict[str,
Any]) -> None:
messageParameters={},
)
+ # Opt-in reuse of a persistent local Spark Connect server. By default
``.remote("local[*]")``
+ # boots a fresh in-process server every process (see
``_start_connect_server`` above); when
+ # reuse is enabled, the first run starts a detached server
(``connect/local_server.py``) and
+ # records it in a discovery file, and later runs reconnect to it instead
of re-paying the cold
+ # start.
+
+ @staticmethod
+ def _local_connect_discovery_path() -> str:
+ """Location of the discovery file describing the running persistent
local server."""
+ override = os.environ.get("SPARK_LOCAL_CONNECT_DISCOVERY")
+ if override:
+ return override
+ return os.path.join(os.path.expanduser("~"), ".spark",
"connect-local.json")
+
+ @staticmethod
+ def _read_local_connect_discovery() -> Optional[Dict[str, Any]]:
+ """Read and validate the discovery file, returning ``None`` if it is
absent or malformed."""
+ path = SparkSession._local_connect_discovery_path()
+ try:
+ with open(path, "r") as f:
+ disc = json.load(f)
+ except (OSError, ValueError):
+ return None
+ if not isinstance(disc, dict) or not all(
+ k in disc for k in ("host", "port", "token", "pid",
"spark_version")
+ ):
+ return None
+ return disc
+
+ @staticmethod
+ def _local_connect_server_is_reusable(disc: Dict[str, Any]) -> bool:
+ """Decide whether the server described by ``disc`` can be reused by
this process.
+
+ Reuse requires that the recorded Spark version matches this client's,
the recorded process
+ is still alive, and it is accepting connections on the recorded port.
A version mismatch,
+ dead pid, or closed port means we must start our own server instead.
+ """
+ import socket
+ from pyspark.version import __version__
+
+ if disc.get("spark_version") != __version__:
+ return False
+ try:
+ os.kill(int(disc["pid"]), 0)
+ except (ProcessLookupError, ValueError, TypeError):
+ return False
+ except OSError:
+ # The process exists but is not ours to signal (e.g.
PermissionError) -- treat as alive.
+ pass
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+ sock.settimeout(0.5)
+ if sock.connect_ex((disc["host"], int(disc["port"]))) != 0:
+ return False
+ return True
+
+ @staticmethod
+ def _reuse_from_discovery() -> Optional[str]:
+ """Return an endpoint for the recorded server if it is reusable, else
``None``.
+
+ On success it also sets ``SPARK_CONNECT_AUTHENTICATE_TOKEN`` so the
client authenticates
+ against that server.
+ """
+ disc = SparkSession._read_local_connect_discovery()
+ if disc is not None and
SparkSession._local_connect_server_is_reusable(disc):
+ os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = disc["token"]
+ return "sc://{}:{}".format(disc["host"], disc["port"])
+ return None
+
+ @staticmethod
+ def _acquire_local_connect_start_lock() -> Any:
+ """Take an exclusive file lock guarding persistent-server start-up.
+
+ Returns the open fd to pass to ``_release_local_connect_start_lock``.
Returns ``None`` on
+ platforms without ``fcntl``; those callers may race and then reconnect
to the server that
+ wins the discovery-file update.
+ """
+ try:
+ import fcntl
+ except ImportError:
+ return None
+ path = SparkSession._local_connect_discovery_path() + ".lock"
+ parent = os.path.dirname(path)
+ if parent and not os.path.isdir(parent):
+ os.makedirs(parent, exist_ok=True)
+ fd = os.open(path, os.O_RDWR | os.O_CREAT, 0o600)
+ fcntl.flock(fd, fcntl.LOCK_EX)
+ return fd
+
+ @staticmethod
+ def _release_local_connect_start_lock(fd: Any) -> None:
+ if fd is None:
+ return
+ try:
+ import fcntl
+
+ fcntl.flock(fd, fcntl.LOCK_UN)
+ finally:
+ os.close(fd)
+
+ @staticmethod
+ def _reuse_or_start_local_connect_server(master: str, opts: Dict[str,
Any]) -> str:
+ """Reuse a running persistent local Connect server, or start one if
none is reusable.
+
+ Returns the ``sc://host:port`` endpoint to connect to. This is the
opt-in counterpart of
+ ``_start_connect_server`` and is only reached for a ``local`` master
when
+ ``spark.local.connect.reuse`` / ``SPARK_LOCAL_CONNECT_REUSE`` is set.
+ """
+ # Fast path: reuse an already-running server without taking the
cross-process lock.
+ endpoint = SparkSession._reuse_from_discovery()
+ if endpoint is not None:
+ return endpoint
+ # No reusable server yet. Serialize start-up across processes when
file locking is
+ # available. Without it, racing callers may each start a daemon; the
winner writes the
+ # discovery file and the others reconnect to it.
+ lock_fd = SparkSession._acquire_local_connect_start_lock()
+ try:
+ endpoint = SparkSession._reuse_from_discovery()
+ if endpoint is not None:
+ return endpoint
+ return SparkSession._start_persistent_local_connect_server(master,
opts)
+ finally:
+ SparkSession._release_local_connect_start_lock(lock_fd)
+
+ @staticmethod
+ def _local_port_available(port: int) -> bool:
+ """Whether ``port`` can currently be bound on localhost (best effort,
subject to races)."""
+ import socket
+
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+ try:
+ sock.bind(("localhost", port))
+ return True
+ except OSError:
+ return False
+
+ @staticmethod
+ def _local_connect_server_conf(opts: Dict[str, Any]) -> Dict[str, Any]:
+ """Start-up configs to seed a freshly started persistent server.
+
+ Mirrors the merge that the in-process ``_start_connect_server``
applies to its ``SparkConf``
+ so that first-run behavior matches (warehouse dir, app name,
jars/packages, catalog confs,
+ etc.). Keys the daemon controls itself (master, port, token, plugins)
and the reuse opt-in
+ keys are excluded. This only seeds the run that *starts* the server; a
later run
+ reconnecting to an already-warm JVM cannot change its static configs.
+ """
+ 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)
+ for k in (
+ "spark.remote",
+ "spark.api.mode",
+ "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",
+ ):
+ conf.pop(k, None)
+ return conf
+
+ @staticmethod
+ def _signal_local_connect_server(pid: int, sig: int) -> bool:
+ """Signal a detached local Connect daemon, including its JVM on POSIX.
+
+ The daemon is launched as a session leader, so its process group id
equals its pid and
+ signalling the group reaps its child JVM too. If the group id differs,
``pid`` belongs to
+ an unrelated process (e.g. recycled after a stale discovery file), so
only the pid itself
+ is signalled -- never a group this code did not create.
+ """
+ try:
+ if os.name == "posix" and os.getpgid(pid) == pid:
+ os.killpg(pid, sig)
+ else:
+ os.kill(pid, sig)
+ return True
+ except OSError:
+ return False
+
+ @staticmethod
+ def _terminate_local_connect_server(proc: Any) -> None:
+ """Terminate a daemon started by
``_start_persistent_local_connect_server`` and its JVM.
+
+ The group SIGTERM asks the daemon and its JVM to shut down gracefully,
but the daemon
+ exiting does not mean the JVM is gone: its graceful shutdown can
outlive the daemon. On
+ POSIX this therefore waits for the whole process group to disappear
and escalates to a
+ group SIGKILL if any member outlives the grace period.
+ """
+ if not SparkSession._signal_local_connect_server(proc.pid,
signal.SIGTERM):
+ return
+ try:
+ proc.wait(timeout=10)
+ except Exception:
+ pass
+ if os.name != "posix":
+ if proc.poll() is None:
+ proc.kill()
+ return
+ # The group id stays valid while any member (i.e. the JVM) is alive,
even after the
+ # daemon leader has been reaped, so poll and signal the group id
directly rather than
+ # via _signal_local_connect_server (whose getpgid lookup needs a live
leader). Signalling
+ # this group is safe: it was created above us by this launch, not read
from disk.
+ deadline = time.time() + 10
+ while time.time() < deadline:
+ proc.poll() # reap the daemon if it exited, so only live members
keep the group
+ try:
+ os.killpg(proc.pid, 0)
+ except OSError:
+ return
+ time.sleep(0.2)
+ try:
+ os.killpg(proc.pid, signal.SIGKILL)
+ except OSError:
+ pass
+ proc.poll()
Review Comment:
It reaped the SIGKILLed daemon so there's no zombie until the client exits.
Replaced with a short timeout explaining why.
##########
python/pyspark/sql/connect/session.py:
##########
@@ -1247,6 +1249,372 @@ def _start_connect_server(master: str, opts: Dict[str,
Any]) -> None:
messageParameters={},
)
+ # Opt-in reuse of a persistent local Spark Connect server. By default
``.remote("local[*]")``
+ # boots a fresh in-process server every process (see
``_start_connect_server`` above); when
+ # reuse is enabled, the first run starts a detached server
(``connect/local_server.py``) and
+ # records it in a discovery file, and later runs reconnect to it instead
of re-paying the cold
+ # start.
+
+ @staticmethod
+ def _local_connect_discovery_path() -> str:
+ """Location of the discovery file describing the running persistent
local server."""
+ override = os.environ.get("SPARK_LOCAL_CONNECT_DISCOVERY")
+ if override:
+ return override
+ return os.path.join(os.path.expanduser("~"), ".spark",
"connect-local.json")
+
+ @staticmethod
+ def _read_local_connect_discovery() -> Optional[Dict[str, Any]]:
+ """Read and validate the discovery file, returning ``None`` if it is
absent or malformed."""
+ path = SparkSession._local_connect_discovery_path()
+ try:
+ with open(path, "r") as f:
+ disc = json.load(f)
+ except (OSError, ValueError):
+ return None
+ if not isinstance(disc, dict) or not all(
+ k in disc for k in ("host", "port", "token", "pid",
"spark_version")
+ ):
+ return None
+ return disc
+
+ @staticmethod
+ def _local_connect_server_is_reusable(disc: Dict[str, Any]) -> bool:
+ """Decide whether the server described by ``disc`` can be reused by
this process.
+
+ Reuse requires that the recorded Spark version matches this client's,
the recorded process
+ is still alive, and it is accepting connections on the recorded port.
A version mismatch,
+ dead pid, or closed port means we must start our own server instead.
+ """
+ import socket
+ from pyspark.version import __version__
+
+ if disc.get("spark_version") != __version__:
+ return False
+ try:
+ os.kill(int(disc["pid"]), 0)
+ except (ProcessLookupError, ValueError, TypeError):
+ return False
+ except OSError:
+ # The process exists but is not ours to signal (e.g.
PermissionError) -- treat as alive.
+ pass
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+ sock.settimeout(0.5)
+ if sock.connect_ex((disc["host"], int(disc["port"]))) != 0:
+ return False
+ return True
+
+ @staticmethod
+ def _reuse_from_discovery() -> Optional[str]:
+ """Return an endpoint for the recorded server if it is reusable, else
``None``.
+
+ On success it also sets ``SPARK_CONNECT_AUTHENTICATE_TOKEN`` so the
client authenticates
+ against that server.
+ """
+ disc = SparkSession._read_local_connect_discovery()
+ if disc is not None and
SparkSession._local_connect_server_is_reusable(disc):
+ os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = disc["token"]
+ return "sc://{}:{}".format(disc["host"], disc["port"])
+ return None
+
+ @staticmethod
+ def _acquire_local_connect_start_lock() -> Any:
+ """Take an exclusive file lock guarding persistent-server start-up.
+
+ Returns the open fd to pass to ``_release_local_connect_start_lock``.
Returns ``None`` on
+ platforms without ``fcntl``; those callers may race and then reconnect
to the server that
+ wins the discovery-file update.
+ """
+ try:
+ import fcntl
+ except ImportError:
+ return None
+ path = SparkSession._local_connect_discovery_path() + ".lock"
+ parent = os.path.dirname(path)
+ if parent and not os.path.isdir(parent):
+ os.makedirs(parent, exist_ok=True)
+ fd = os.open(path, os.O_RDWR | os.O_CREAT, 0o600)
+ fcntl.flock(fd, fcntl.LOCK_EX)
+ return fd
+
+ @staticmethod
+ def _release_local_connect_start_lock(fd: Any) -> None:
+ if fd is None:
+ return
+ try:
+ import fcntl
+
+ fcntl.flock(fd, fcntl.LOCK_UN)
+ finally:
+ os.close(fd)
+
+ @staticmethod
+ def _reuse_or_start_local_connect_server(master: str, opts: Dict[str,
Any]) -> str:
+ """Reuse a running persistent local Connect server, or start one if
none is reusable.
+
+ Returns the ``sc://host:port`` endpoint to connect to. This is the
opt-in counterpart of
+ ``_start_connect_server`` and is only reached for a ``local`` master
when
+ ``spark.local.connect.reuse`` / ``SPARK_LOCAL_CONNECT_REUSE`` is set.
+ """
+ # Fast path: reuse an already-running server without taking the
cross-process lock.
+ endpoint = SparkSession._reuse_from_discovery()
+ if endpoint is not None:
+ return endpoint
+ # No reusable server yet. Serialize start-up across processes when
file locking is
+ # available. Without it, racing callers may each start a daemon; the
winner writes the
+ # discovery file and the others reconnect to it.
+ lock_fd = SparkSession._acquire_local_connect_start_lock()
+ try:
+ endpoint = SparkSession._reuse_from_discovery()
+ if endpoint is not None:
+ return endpoint
+ return SparkSession._start_persistent_local_connect_server(master,
opts)
+ finally:
+ SparkSession._release_local_connect_start_lock(lock_fd)
+
+ @staticmethod
+ def _local_port_available(port: int) -> bool:
+ """Whether ``port`` can currently be bound on localhost (best effort,
subject to races)."""
+ import socket
+
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+ try:
+ sock.bind(("localhost", port))
+ return True
+ except OSError:
+ return False
+
+ @staticmethod
+ def _local_connect_server_conf(opts: Dict[str, Any]) -> Dict[str, Any]:
+ """Start-up configs to seed a freshly started persistent server.
+
+ Mirrors the merge that the in-process ``_start_connect_server``
applies to its ``SparkConf``
+ so that first-run behavior matches (warehouse dir, app name,
jars/packages, catalog confs,
+ etc.). Keys the daemon controls itself (master, port, token, plugins)
and the reuse opt-in
+ keys are excluded. This only seeds the run that *starts* the server; a
later run
+ reconnecting to an already-warm JVM cannot change its static configs.
+ """
+ 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)
+ for k in (
+ "spark.remote",
+ "spark.api.mode",
+ "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",
+ ):
+ conf.pop(k, None)
+ return conf
+
+ @staticmethod
+ def _signal_local_connect_server(pid: int, sig: int) -> bool:
+ """Signal a detached local Connect daemon, including its JVM on POSIX.
+
+ The daemon is launched as a session leader, so its process group id
equals its pid and
+ signalling the group reaps its child JVM too. If the group id differs,
``pid`` belongs to
+ an unrelated process (e.g. recycled after a stale discovery file), so
only the pid itself
+ is signalled -- never a group this code did not create.
+ """
+ try:
+ if os.name == "posix" and os.getpgid(pid) == pid:
+ os.killpg(pid, sig)
+ else:
+ os.kill(pid, sig)
+ return True
+ except OSError:
+ return False
+
+ @staticmethod
+ def _terminate_local_connect_server(proc: Any) -> None:
+ """Terminate a daemon started by
``_start_persistent_local_connect_server`` and its JVM.
+
+ The group SIGTERM asks the daemon and its JVM to shut down gracefully,
but the daemon
+ exiting does not mean the JVM is gone: its graceful shutdown can
outlive the daemon. On
+ POSIX this therefore waits for the whole process group to disappear
and escalates to a
+ group SIGKILL if any member outlives the grace period.
+ """
+ if not SparkSession._signal_local_connect_server(proc.pid,
signal.SIGTERM):
+ return
+ try:
+ proc.wait(timeout=10)
+ except Exception:
+ pass
+ if os.name != "posix":
+ if proc.poll() is None:
+ proc.kill()
+ return
+ # The group id stays valid while any member (i.e. the JVM) is alive,
even after the
+ # daemon leader has been reaped, so poll and signal the group id
directly rather than
+ # via _signal_local_connect_server (whose getpgid lookup needs a live
leader). Signalling
+ # this group is safe: it was created above us by this launch, not read
from disk.
+ deadline = time.time() + 10
+ while time.time() < deadline:
+ proc.poll() # reap the daemon if it exited, so only live members
keep the group
+ try:
+ os.killpg(proc.pid, 0)
+ except OSError:
+ return
+ time.sleep(0.2)
+ try:
+ os.killpg(proc.pid, signal.SIGKILL)
+ except OSError:
+ pass
+ proc.poll()
Review Comment:
It reaped the SIGKILLed daemon so there's no zombie until the client exits.
Replaced with a short timeout and comment explaining why.
--
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]
