From 8c4815098b7732eb377332c18b804e2576400c7b Mon Sep 17 00:00:00 2001 From: Lisa Cao Date: Wed, 6 May 2026 10:22:50 -0700 Subject: [PATCH] Fix SparkPipelinesHook for Spark Connect (sc://) connections The Connect-native spark-pipelines CLI in Spark 4.x rejects the cluster- manager flags (--master, --deploy-mode, --name) that SparkSubmitHook unconditionally emits, with `SparkException: Remote cannot be specified with master and/or deploy mode`. The bash launcher itself also starts a JVM SparkContext that tries to bind a Connect server on the same port the long-running Connect daemon is already on. When the configured conn_id resolves to a spark_connect-typed connection (handled by the existing SparkConnectHook), this hook now: * Skips _build_spark_common_args() so cluster-manager flags are not emitted. * Sets SPARK_REMOTE in the subprocess env from SparkConnectHook.get_connection_url(), letting the operator caller override it via env_vars for failover scenarios. * Invokes pyspark.pipelines.cli as a Python module instead of the spark-pipelines bash launcher, so SparkSession.builder.getOrCreate() becomes a Connect client rather than starting an in-process JVM SparkContext that conflicts with the daemon. The spark-submit-style code path (spark / yarn / k8s connection types) is unchanged. Connection-lookup failures fall through to the legacy path so existing unit tests with no metadata DB continue to pass. Adds three tests for the spark_connect path covering (a) command shape without cluster-manager flags, (b) SPARK_REMOTE injection from the connection URI, and (c) caller-supplied SPARK_REMOTE preservation for failover. --- .../apache/spark/hooks/spark_pipelines.py | 96 +++++++++++-- .../spark/hooks/test_spark_pipelines.py | 130 ++++++++++++++++++ 2 files changed, 215 insertions(+), 11 deletions(-) diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_pipelines.py b/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_pipelines.py index e721433b93b39..4d2457668a77e 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_pipelines.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_pipelines.py @@ -17,11 +17,13 @@ # under the License. from __future__ import annotations +import os import subprocess +import sys from typing import Any -from airflow.providers.apache.spark.hooks.spark_submit import SparkSubmitHook -from airflow.providers.common.compat.sdk import AirflowException +from airflow.providers.apache.spark.hooks.spark_submit import DEFAULT_SPARK_BINARY, SparkSubmitHook +from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException class SparkPipelinesException(AirflowException): @@ -35,6 +37,20 @@ class SparkPipelinesHook(SparkSubmitHook): Extends SparkSubmitHook to leverage existing connection management while providing pipeline-specific functionality. + Two connection modes are supported: + + * Legacy spark-submit-style (``spark`` / ``yarn`` / ``k8s`` connection types) — + invokes the ``spark-pipelines`` launcher with ``--master``, ``--deploy-mode`` + and the rest of the standard cluster-manager flags assembled by + :class:`~airflow.providers.apache.spark.hooks.spark_submit.SparkSubmitHook`. + * Spark Connect (``spark_connect`` connection type, Spark 4.x+) — sets + ``SPARK_REMOTE`` from the connection's ``sc://`` URI and invokes the + Connect-native ``pyspark.pipelines.cli`` Python module directly. The + cluster-manager flags are *not* emitted: the Connect-native CLI rejects + them with ``SparkException: Remote cannot be specified with master and/or + deploy mode``, and the ``spark-pipelines`` bash launcher itself starts a + JVM ``SparkContext`` that collides with the Connect daemon's gRPC port. + :param pipeline_spec: Path to the pipeline specification file (YAML) :param pipeline_command: The spark-pipelines command to run ('run', 'dry-run') """ @@ -52,7 +68,56 @@ def __init__( if pipeline_command not in ["run", "dry-run"]: raise ValueError(f"Invalid pipeline command: {pipeline_command}. Must be 'run' or 'dry-run'") + def _resolve_connection(self) -> dict[str, Any]: + """ + Resolve the configured connection, branching on Spark Connect. + + For ``spark_connect``-typed connections, populate a ``spark_remote`` key + with the ``sc://`` URI rendered by + :class:`~airflow.providers.apache.spark.hooks.spark_connect.SparkConnectHook` + and zero out the spark-submit cluster-manager fields — the Connect-native + CLI doesn't consume them. For every other connection type, defer to the + parent's resolver so spark-submit-style behaviour is unchanged. + """ + try: + conn = self.get_connection(self._conn_id) + except AirflowNotFoundException: + # No connection configured — fall through to spark-submit defaults. + return super()._resolve_connection() + + if conn.conn_type != "spark_connect": + return super()._resolve_connection() + + # Local import: SparkConnectHook lives in the same provider but loading + # it eagerly would create an unnecessary cycle for spark-submit-only + # deployments that never touch Connect. + from airflow.providers.apache.spark.hooks.spark_connect import SparkConnectHook + + return { + # ``master`` is consumed by the parent's ``__init__`` (substring + # checks for ``yarn``/``k8s``/``spark://``); leave it empty so none + # of those branches match. Connect mode never calls + # ``_build_spark_common_args``, which is the only thing that would + # actually emit ``--master`` to the CLI. + "master": "", + "queue": None, + "deploy_mode": None, + "spark_binary": self.spark_binary or DEFAULT_SPARK_BINARY, + "namespace": None, + "principal": self._principal, + "keytab": self._keytab, + "spark_remote": SparkConnectHook(conn_id=self._conn_id).get_connection_url(), + } + def _get_spark_binary_path(self) -> list[str]: + if self._connection.get("spark_remote"): + # The ``spark-pipelines`` bash launcher routes through + # ``spark-class`` → ``SparkSubmit`` → JVM ``SparkContext``, which + # appends cluster-manager flags the Connect-native CLI rejects and + # binds an in-process Connect server that collides with the + # long-running daemon. Invoke the underlying Python module directly + # so ``SparkSession.builder.getOrCreate()`` becomes a Connect client. + return [sys.executable, "-m", "pyspark.pipelines.cli"] return ["spark-pipelines"] def _build_spark_pipelines_command(self) -> list[str]: @@ -61,16 +126,17 @@ def _build_spark_pipelines_command(self) -> list[str]: :return: full command to be executed """ - # Start with spark-pipelines binary and command connection_cmd = self._get_spark_binary_path() connection_cmd.append(self.pipeline_command) - # Add pipeline spec if provided if self.pipeline_spec: connection_cmd.extend(["--spec", self.pipeline_spec]) - # Reuse parent's common spark argument building logic - connection_cmd.extend(self._build_spark_common_args()) + if not self._connection.get("spark_remote"): + # Reuse parent's common spark argument building logic. The + # Connect-native CLI rejects --master/--deploy-mode/--name, so only + # emit them for legacy spark-submit-style connections. + connection_cmd.extend(self._build_spark_common_args()) self.log.info("Spark-Pipelines cmd: %s", self._mask_cmd(connection_cmd)) return connection_cmd @@ -82,12 +148,20 @@ def submit_pipeline(self, **kwargs: Any) -> None: :param kwargs: extra arguments to Popen (see subprocess.Popen) """ pipelines_cmd = self._build_spark_pipelines_command() - - if self._env: - import os - + spark_remote = self._connection.get("spark_remote") + + # ``self._env`` is only populated by ``_build_spark_common_args`` — + # which the Connect path skips — so fall back to ``self._env_vars`` + # (the operator-supplied env_vars kwarg) for that path. + env_overrides: dict[str, str] = dict(self._env or self._env_vars or {}) + if spark_remote: + # Don't clobber a SPARK_REMOTE the operator caller already set via + # env_vars; that takes precedence for failover routing. + env_overrides.setdefault("SPARK_REMOTE", spark_remote) + + if env_overrides: env = os.environ.copy() - env.update(self._env) + env.update(env_overrides) kwargs["env"] = env self._submit_sp = subprocess.Popen( diff --git a/providers/apache/spark/tests/unit/apache/spark/hooks/test_spark_pipelines.py b/providers/apache/spark/tests/unit/apache/spark/hooks/test_spark_pipelines.py index abf1d5d94ec2b..6c29399c12ea0 100644 --- a/providers/apache/spark/tests/unit/apache/spark/hooks/test_spark_pipelines.py +++ b/providers/apache/spark/tests/unit/apache/spark/hooks/test_spark_pipelines.py @@ -17,6 +17,7 @@ # under the License. from __future__ import annotations +import sys from unittest.mock import MagicMock, patch import pytest @@ -183,3 +184,132 @@ def test_submit_calls_submit_pipeline(self): with patch.object(self.hook, "submit_pipeline") as mock_submit_pipeline: self.hook.submit("dummy_application") mock_submit_pipeline.assert_called_once() + + # ------------------------------------------------------------------ # + # Spark Connect path: spark_connect-typed conn_id should bypass the # + # spark-submit launcher and the cluster-manager flags entirely, and # + # set SPARK_REMOTE in the subprocess environment instead. # + # ------------------------------------------------------------------ # + + @patch( + "airflow.providers.apache.spark.hooks.spark_connect.SparkConnectHook.get_connection_url", + return_value="sc://spark-connect.example:15002/", + ) + @patch("airflow.providers.apache.spark.hooks.spark_pipelines.SparkPipelinesHook.get_connection") + def test_build_pipelines_command_spark_connect_skips_cluster_args( + self, mock_get_connection, mock_get_url + ): + spark_connect_conn = MagicMock() + spark_connect_conn.conn_type = "spark_connect" + spark_connect_conn.host = "spark-connect.example" + spark_connect_conn.port = 15002 + spark_connect_conn.login = None + spark_connect_conn.password = None + # SparkSubmitHook.__init__ -> _resolve_connection() reads + # extra_dejson.get("queue"|"deploy-mode"|"spark-binary"|"namespace"); + # an empty dict makes it default-and-skip cleanly. + spark_connect_conn.extra_dejson = {} + mock_get_connection.return_value = spark_connect_conn + + hook = SparkPipelinesHook( + pipeline_spec="test_pipeline.yml", + pipeline_command="run", + conn_id="spark_connect_default", + num_executors=2, + executor_memory="2G", + deploy_mode="client", + ) + + cmd = hook._build_spark_pipelines_command() + + # Connect-native CLI is invoked via the python module — bypassing the + # bash launcher that would otherwise start a colliding JVM Connect server. + assert cmd[:3] == [sys.executable, "-m", "pyspark.pipelines.cli"] + assert "run" in cmd + assert "--spec" in cmd + assert "test_pipeline.yml" in cmd + # Cluster-manager args MUST NOT appear: the Connect-native CLI rejects + # them with `Remote cannot be specified with master and/or deploy mode`. + assert "--master" not in cmd + assert "--deploy-mode" not in cmd + assert "--name" not in cmd + assert "--num-executors" not in cmd + assert "--executor-memory" not in cmd + + @patch("subprocess.Popen") + @patch( + "airflow.providers.apache.spark.hooks.spark_connect.SparkConnectHook.get_connection_url", + return_value="sc://spark-connect.example:15002/", + ) + @patch("airflow.providers.apache.spark.hooks.spark_pipelines.SparkPipelinesHook.get_connection") + def test_submit_pipeline_spark_connect_sets_spark_remote( + self, mock_get_connection, mock_get_url, mock_popen + ): + spark_connect_conn = MagicMock() + spark_connect_conn.conn_type = "spark_connect" + spark_connect_conn.host = "spark-connect.example" + spark_connect_conn.port = 15002 + spark_connect_conn.login = None + spark_connect_conn.password = None + # SparkSubmitHook.__init__ -> _resolve_connection() reads + # extra_dejson.get("queue"|"deploy-mode"|"spark-binary"|"namespace"); + # an empty dict makes it default-and-skip cleanly. + spark_connect_conn.extra_dejson = {} + mock_get_connection.return_value = spark_connect_conn + + mock_process = MagicMock() + mock_process.wait.return_value = 0 + mock_process.stdout = ["Run is COMPLETED."] + mock_popen.return_value = mock_process + + hook = SparkPipelinesHook( + pipeline_spec="test_pipeline.yml", + pipeline_command="run", + conn_id="spark_connect_default", + ) + hook.submit_pipeline() + + mock_popen.assert_called_once() + _, popen_kwargs = mock_popen.call_args + env = popen_kwargs["env"] + assert env["SPARK_REMOTE"] == "sc://spark-connect.example:15002/" + + @patch("subprocess.Popen") + @patch( + "airflow.providers.apache.spark.hooks.spark_connect.SparkConnectHook.get_connection_url", + return_value="sc://spark-connect.example:15002/", + ) + @patch("airflow.providers.apache.spark.hooks.spark_pipelines.SparkPipelinesHook.get_connection") + def test_submit_pipeline_spark_connect_preserves_caller_spark_remote( + self, mock_get_connection, mock_get_url, mock_popen + ): + # If the operator caller has already set SPARK_REMOTE via env_vars + # (e.g. routing to a different daemon for failover), the hook must + # not clobber it with the connection's URI. + spark_connect_conn = MagicMock() + spark_connect_conn.conn_type = "spark_connect" + spark_connect_conn.host = "spark-connect.example" + spark_connect_conn.port = 15002 + spark_connect_conn.login = None + spark_connect_conn.password = None + # SparkSubmitHook.__init__ -> _resolve_connection() reads + # extra_dejson.get("queue"|"deploy-mode"|"spark-binary"|"namespace"); + # an empty dict makes it default-and-skip cleanly. + spark_connect_conn.extra_dejson = {} + mock_get_connection.return_value = spark_connect_conn + + mock_process = MagicMock() + mock_process.wait.return_value = 0 + mock_process.stdout = [] + mock_popen.return_value = mock_process + + hook = SparkPipelinesHook( + pipeline_spec="test_pipeline.yml", + pipeline_command="run", + conn_id="spark_connect_default", + env_vars={"SPARK_REMOTE": "sc://override.example:15002/"}, + ) + hook.submit_pipeline() + + _, popen_kwargs = mock_popen.call_args + assert popen_kwargs["env"]["SPARK_REMOTE"] == "sc://override.example:15002/"