ericm-db commented on code in PR #56907:
URL: https://github.com/apache/spark/pull/56907#discussion_r3540188980
##########
python/pyspark/sql/connect/session.py:
##########
@@ -1247,6 +1249,347 @@ 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``,
or ``None`` when file
+ locking is unavailable (e.g. Windows, which has no ``fcntl``); there
we rely on the
+ reconnect-the-winner fallback in
``_start_persistent_local_connect_server`` instead.
+ """
+ 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 so
concurrent opted-in
+ # processes (parallel workers, an IDE spawning scripts) do not each
spawn a server and
+ # collide on the port; the winner writes the discovery file and the
others reuse 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 _terminate_local_connect_server(proc: Any) -> None:
+ """Terminate a daemon started by
``_start_persistent_local_connect_server`` and its JVM.
+
+ The daemon is a POSIX session leader (``start_new_session=True``) and
its child JVM stays in
+ that process group, so signalling the group reaps both -- important
when the timeout fires
+ before the daemon has wired up its own signal handling. Escalates to
SIGKILL if a graceful
+ stop does not take. On non-POSIX platforms it falls back to
terminating just the process.
+ """
+ try:
+ if os.name == "posix":
+ os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
+ else:
+ proc.terminate()
+ try:
+ proc.wait(timeout=10)
+ except Exception:
+ if os.name == "posix":
+ os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
+ else:
+ proc.kill()
+ except OSError:
+ pass
+
+ @staticmethod
+ def _start_persistent_local_connect_server(master: str, opts: Dict[str,
Any]) -> str:
+ """Launch a detached persistent local Connect server and wait until it
is reachable.
+
+ Callers must hold the start-up lock (see
``_reuse_or_start_local_connect_server``). If the
+ server cannot be started but another process has meanwhile published a
reusable one, this
+ reconnects to it rather than failing.
+ """
+ import socket
+ import subprocess
+ import tempfile
+
+ discovery_path = SparkSession._local_connect_discovery_path()
+ token = opts.get("spark.connect.authenticate.token") or
str(uuid.uuid4())
+
+ # Choose the port. Tests use an ephemeral port (0) so they can run in
parallel. Otherwise we
+ # honor the configured/default port, but fall back to an ephemeral one
if it is already
+ # taken -- e.g. by a stale server we just rejected on version mismatch
-- so a fresh server
+ # can still start instead of failing to bind.
+ if "SPARK_TESTING" in os.environ:
+ port = 0
+ else:
+ port = int(
+ opts.get("spark.local.connect.server.port",
DefaultChannelBuilder.default_port())
+ )
+ if port != 0 and not SparkSession._local_port_available(port):
+ port = 0
+ idle_timeout = opts.get("spark.local.connect.server.idleTimeout",
"3600")
+
+ # Seed the server with the caller's start-up confs (warehouse dir,
jars, catalog, etc.) so
+ # first-run behavior matches the in-process path. Passed as a JSON
file since confs are
+ # arbitrary key/values. Written under the discovery dir with 0600
perms as it may hold
+ # sensitive values, and removed once the daemon has started.
+ conf_file = None
+ seed_conf = SparkSession._local_connect_server_conf(opts)
+ if seed_conf:
+ parent = os.path.dirname(discovery_path)
+ if parent and not os.path.isdir(parent):
+ os.makedirs(parent, exist_ok=True)
+ fd, conf_file = tempfile.mkstemp(prefix="connect-local-conf-",
dir=parent or None)
+ with os.fdopen(fd, "w") as f:
+ json.dump(seed_conf, f)
+ os.chmod(conf_file, 0o600)
+
+ daemon = os.path.join(os.path.dirname(os.path.abspath(__file__)),
"local_server.py")
+ cmd = [
+ sys.executable,
+ daemon,
+ "--master",
+ master,
+ "--port",
+ str(port),
+ "--token",
+ token,
+ "--discovery",
+ discovery_path,
+ "--idle-timeout",
+ str(idle_timeout),
+ ]
+ if conf_file is not None:
+ cmd += ["--conf-file", conf_file]
+
+ # Launch detached so the server outlives this client process.
+ env = dict(os.environ)
+ for var in ("SPARK_REMOTE", "SPARK_LOCAL_REMOTE",
"SPARK_CONNECT_MODE_ENABLED"):
+ env.pop(var, None)
+ popen_kwargs: Dict[str, Any] = dict(
+ env=env,
+ stdin=subprocess.DEVNULL,
+ stdout=subprocess.DEVNULL,
+ stderr=subprocess.DEVNULL,
+ )
+ if os.name == "posix":
+ popen_kwargs["start_new_session"] = True
+ else:
+ detached = getattr(subprocess, "DETACHED_PROCESS", 0)
+ new_group = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)
+ popen_kwargs["creationflags"] = detached | new_group
+ proc = subprocess.Popen(cmd, **popen_kwargs)
+
+ try:
+ # Wait for the server to write the discovery file (which records
the actual bound port)
+ # and start accepting connections.
+ deadline = time.time() + 120
+ while time.time() < deadline:
+ exit_code = proc.poll()
+ if exit_code is not None:
+ # Our daemon died. Another process may have published a
usable server in the
+ # meantime (e.g. a port race), so prefer reconnecting to
it over failing.
+ endpoint = SparkSession._reuse_from_discovery()
+ if endpoint is not None:
+ return endpoint
+ raise PySparkRuntimeError(
+ errorClass="LOCAL_CONNECT_SERVER_START_FAILED",
+ messageParameters={
+ "reason": "the server process exited with code
{}".format(exit_code)
+ },
+ )
+ disc = SparkSession._read_local_connect_discovery()
+ if disc is not None and disc.get("pid") == proc.pid and
disc.get("token") == token:
+ 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:
+ os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] =
token
+ return "sc://{}:{}".format(disc["host"],
disc["port"])
+ time.sleep(0.25)
+
+ # Timed out. The daemon may still be inside getOrCreate() -- i.e.
before it has wired up
+ # its own SIGTERM handler -- and it has already spawned a child
JVM. Signal the whole
+ # process group (the daemon is a session leader via
start_new_session) and escalate to
+ # SIGKILL, so the JVM is reaped rather than orphaned.
+ SparkSession._terminate_local_connect_server(proc)
+ raise PySparkRuntimeError(
+ errorClass="LOCAL_CONNECT_SERVER_START_FAILED",
+ messageParameters={"reason": "the server did not become ready
within 120 seconds"},
+ )
+ finally:
+ if conf_file is not None:
+ try:
+ os.remove(conf_file)
+ except OSError:
+ pass
+
+ @staticmethod
+ def _stop_local_connect_server() -> bool:
Review Comment:
Right, that makes sense. done
--
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]
